Remove interpretability directory and update README guides

Remove the 31 files under BioMod/interpretability as requested. Update the two README introductions, workflow references, checkpoint labels and directory trees. No external code or notebook imports depend on the removed directory. Model weights, paired configuration, raw data and all other repository files are unchanged.

BioMod/README.md

@@ -27,7 +27,7 @@ This is the `BioMod/` source directory inside the project root. See the
 from these repository trees.
 
 ```text
-BioMod/                                                   # Task code: training, evaluation, interpretability and tutorials
+BioMod/                                                   # Task code: training, evaluation and tutorials
 ├── cell_annotation/                                      # Cell-type annotation and model fine-tuning
 │   ├── biomod_data/                                      # Hallmark and Reactome modules with loading utilities
 │   │   ├── load_biomod.py
@@ -51,41 +51,6 @@ BioMod/                                                   # Task code: training,
 │   ├── scfm_utils_local.py
 │   ├── train_biomod.py
 │   └── utils.py
-├── interpretability/                                     # Module representations and biological interpretability
-│   ├── attention/                                        # Module crosstalk and within-module gene attention
-│   │   ├── __init__.py
-│   │   ├── biomod_crosstalk.py
-│   │   ├── biomod_gene_attention.py
-│   │   ├── README.md
-│   │   └── README_gene_attention.md
-│   ├── biomod_importance/                                # Module importance, ablation studies and marker-gene validation
-│   │   ├── __init__.py
-│   │   ├── cross_class_similarity.py
-│   │   ├── family_ablation.py
-│   │   ├── family_importance_analysis.py
-│   │   ├── marker_gene_verification.py
-│   │   ├── masking_ablation.py
-│   │   ├── biomod_importance.py
-│   │   ├── biomod_importance_publication.py
-│   │   ├── pub_utils.py
-│   │   ├── random_baseline.py
-│   │   ├── README.md
-│   │   ├── RESULTS.md
-│   │   ├── tcell_reactome_analysis.py
-│   │   └── TODO.md
-│   ├── biomod_text_alignment/                            # Semantic alignment of module representations and pathway text
-│   │   ├── __init__.py
-│   │   ├── DESIGN.md
-│   │   ├── fetch_reactome_descriptions.py
-│   │   ├── MANUSCRIPT_SECTION.md
-│   │   ├── biomod_text_alignment.py
-│   │   ├── biomod_text_alignment_extended.py
-│   │   ├── biomod_text_alignment_publication.py
-│   │   └── reactome_descriptions.json
-│   ├── __init__.py
-│   ├── biomod_geometry.py
-│   ├── README.md
-│   └── utils.py
 ├── tests/                                                # CPU tests for shared training utilities
 │   └── test_training_utils.py
 ├── toturials/                                            # Jupyter tutorials for five downstream tasks

BioMod/interpretability/README.md

@@ -1,361 +0,0 @@
-# Module Interpretability Experiments for BioMod
-
-This directory contains analyses of BioMod module representations, covering
-latent geometry, pathway-text alignment, module importance, and attention.
-These workflows examine structural, semantic, and functional properties of the
-learned modules. Use the
-[BioMod model weights and matching configuration](../../README.md#biomod-model-weights)
-when loading the pretrained backbone.
-
-## Experiment Roadmap, Ordered by Priority
-
-| # | Script | Core question | Status |
-| --- | --- | --- | --- |
-| 1 | `biomod_geometry.py` | Does the module latent space self-organize by pathway family,  | Completed |
-| 2 | `biomod_text_alignment/biomod_text_alignment.py` | Do module embeddings align with natural-language pathway descriptions,  | Implemented |
-| 3 | `biomod_importance/biomod_importance.py` | Does task-driven module importance match known pathway-cell-type or pathway-disease associations,  | Implemented |
-| 4 | `module_vs_gene.py` | Is BioMod better than GenePT/scGPT for pathway-level interpretability,  | TODO |
-| 6 | `attention/biomod_crosstalk.py` | Do module interactions match Reactome/KEGG pathway relationships,  | Implemented |
-
-Experiments 1 and 2 are the core differentiators because gene-level baselines
-cannot structurally reproduce them. Experiment 3 provides task-level functional
-validation. Experiment 4 provides direct comparison to baselines. Additional
-experiments, such as module crosstalk vs Reactome, module dynamics along
-developmental trajectories, and retraining-based module-granularity ablation, can
-be pursued after the core experiments are complete.
-
-## Experiment 1: Intrinsic Geometry of Module Embeddings (`biomod_geometry.py`)
-
-**Hypothesis:** the BioMod encoder is trained with a masked-module
-reconstruction task, but it should still encode semantically related modules,
-such as modules from the same pathway family or MSigDB hallmark biological
-process, into nearby regions of latent space, even though family labels are
-never seen during training.
-
-**Workflow:**
-
-1. Load the pretrained backbone and run inference with `mask_ratio=0.0` and `return_encodings=True` to obtain module encodings shaped `[B, num_modules, base_dim]`.
-2. Average over the cell population to obtain `[num_modules, base_dim]`, using streaming accumulation rather than storing the full table in memory.
-3. Parse family assignments from the keys in `merged_biomods.json`: `Hallmarks` (50), `Reactome` (491), `Hotspot` (73), and `Processed` (46). Hallmark modules are also annotated with one of eight MSigDB biological-process categories: immune, metabolic, signaling, proliferation, development, dna_damage, pathway, and cellular_component.
-4. Quantify family separation:
-   - silhouette coefficient, using both cosine and Euclidean metrics.
-   - kNN purity at `k=10`, using both cosine and Euclidean metrics.
-   - ratio of mean inter-family to intra-family distance.
-5. Compare to a permutation null distribution by shuffling family labels `K=200` times, recomputing metrics, and reporting the observed value, null mean, null standard deviation, and right-tail empirical p-value.
-6. Visualize with UMAP and PCA in 2D, colored by family or hallmark biological process.
-
-**Differentiation from gene-level baselines:** gene-level foundation models embed
-single genes rather than pathway groups. To reproduce this experiment, they would
-first have to manually aggregate gene embeddings by pathway, which compresses
-gene-level signals and introduces a hand-designed step. BioMod learns this
-aggregation end to end during pretraining.
-
-## File Structure
-
-```text
-BioMod/interpretability/                                 # Module representations and biological interpretability
-├── attention/                                           # Module crosstalk and within-module gene attention
-│   ├── __init__.py
-│   ├── biomod_crosstalk.py
-│   ├── biomod_gene_attention.py
-│   ├── README.md
-│   └── README_gene_attention.md
-├── biomod_importance/                                   # Module importance, ablation studies and marker-gene validation
-│   ├── __init__.py
-│   ├── cross_class_similarity.py
-│   ├── family_ablation.py
-│   ├── family_importance_analysis.py
-│   ├── marker_gene_verification.py
-│   ├── masking_ablation.py
-│   ├── biomod_importance.py
-│   ├── biomod_importance_publication.py
-│   ├── pub_utils.py
-│   ├── random_baseline.py
-│   ├── README.md
-│   ├── RESULTS.md
-│   ├── tcell_reactome_analysis.py
-│   └── TODO.md
-├── biomod_text_alignment/                               # Semantic alignment of module representations and pathway text
-│   ├── __init__.py
-│   ├── DESIGN.md
-│   ├── fetch_reactome_descriptions.py
-│   ├── MANUSCRIPT_SECTION.md
-│   ├── biomod_text_alignment.py
-│   ├── biomod_text_alignment_extended.py
-│   ├── biomod_text_alignment_publication.py
-│   └── reactome_descriptions.json
-├── __init__.py
-├── biomod_geometry.py
-├── README.md
-└── utils.py
-```
-
-## Command-Line Usage
-
-Run all commands below from the BioMod repository root. Bundled default assets
-are located relative to the script files, so those defaults do not depend on the
-working directory. Explicit relative CLI paths still resolve from your current
-working directory.
-
-The h5ad examples use the bundled pancreas dataset. Pretraining `.npz` cell files
-are not included: replace `/absolute/path/to/pretraining_cell_npz` with your own
-prepared directory. Generated embeddings and analysis results are also not
-bundled; run the corresponding upstream analysis before using those paths.
-
-```bash
-# General h5ad input: cell annotation, zero-shot, or any annotated dataset.
-python BioMod/interpretability/biomod_geometry.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl           weight/BioMod.pkl \
-    --biomod_data_dir       model_dependencies/biomod_data_all \
-    --gene_list_path       model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path    model_dependencies/gene2ens.csv \
-    --gene2vec_data_path   model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file        model_dependencies/gene2ens.csv \
-    --data_path            data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --data_mode            h5ad \
-    --gene_name_column     auto \
-    --batch_size           640 \
-    --output_dir           outputs/interpretability/geometry/pancreas_scib
-```
-
-For pretraining data, where each cell is an `.npz` file already in the 27,874
-model-gene space:
-
-```bash
-python BioMod/interpretability/biomod_geometry.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl           weight/BioMod.pkl \
-    --biomod_data_dir       model_dependencies/biomod_data_all \
-    --gene_list_path       model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path    model_dependencies/gene2ens.csv \
-    --gene2vec_data_path   model_dependencies/selected_gene2vec_27k.npy \
-    --data_path            /absolute/path/to/pretraining_cell_npz \
-    --data_mode            pretrain_npz \
-    --max_cells            50000 \
-    --batch_size           512 \
-    --output_dir           outputs/interpretability/geometry/pretrain_subset
-```
-
-### Multi-GPU DDP Acceleration
-
-`torchrun` multi-GPU inference is supported. When `WORLD_SIZE > 1`, each GPU
-processes a subset of cells, embeddings are aggregated with `all_reduce`, and
-rank 0 computes metrics and figures. Single-GPU and CPU behavior is unchanged.
-
-```bash
-torchrun --nproc_per_node=2 BioMod/interpretability/biomod_geometry.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl           weight/BioMod.pkl \
-    --biomod_data_dir       model_dependencies/biomod_data_all \
-    --gene_list_path       model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path    model_dependencies/gene2ens.csv \
-    --gene2vec_data_path   model_dependencies/selected_gene2vec_27k.npy \
-    --data_path            /absolute/path/to/pretraining_cell_npz \
-    --data_mode            pretrain_npz \
-    --max_cells            100000 \
-    --batch_size           16 \
-    --output_dir           outputs/interpretability/geometry/pretrain_ddp
-```
-
-### Common Arguments
-
-| Argument | Default | Description |
-| --- | --- | --- |
-| `--max_cells` | None | Number of cells to subsample; the mean is stable from a few thousand cells onward. |
-| `--cell_num` | None | Total cell count in `pretrain_npz` mode, used for sequential names `cell0.npz` to `cellN.npz`. |
-| `--knn_k` | 10 | `k` for kNN purity. |
-| `--null_permutations` | 200 | Number of label permutations for the empirical null distribution. |
-| `--include_other_in_metrics` | Off | By default, modules from the `Other` family are excluded. |
-| `--shuffle_biomod_genes` | Off | Randomly shuffle gene-to-module assignments as a negative control or ablation. Uses the same flag as pretraining and downstream finetuning. |
-| `--umap_n_neighbors` | 15 | UMAP graph-construction parameter. |
-| `--umap_min_dist` | 0.3 | UMAP layout parameter. |
-
-## Output Files
-
-Files are written to `<output_dir>/`:
-
-| File | Content |
-| --- | --- |
-| `module_embeddings.npy` | Cell-averaged module embeddings with shape `[num_modules, base_dim]` |
-| `module_metadata.csv` | Metadata for each module: `module_name`, `family`, and `hallmark_process` |
-| `metrics.json` | Observed metrics, null mean/std, and right-tail p-values for family-level and hallmark-process-level analyses |
-| `pca_family.png` | 2D PCA colored by family |
-| `umap_family.png` | 2D UMAP colored by family; requires `umap-learn` |
-| `pca_hallmark_process.png` | PCA for Hallmark modules only, colored by MSigDB biological process |
-| `umap_hallmark_process.png` | UMAP for Hallmark modules only, colored by MSigDB biological process |
-| `run.log` | Full log with formatted core metrics |
-
-## Recommended Negative Controls
-
-Two control experiments make Experiment 1 more rigorous:
-
-1. **Random module assignment**: rerun with `--shuffle_biomod_genes`. This flag preserves module sizes but randomly assigns genes to modules, matching the same ablation flag used for pretraining and finetuning. If geometry is pathway-driven rather than a generic property of the encoder, family separation should drop close to the null distribution.
-
-2. **Untrained backbone**: set `--biomod_checkpoint` to the `*_init.pth` checkpoint generated during pretraining with `--save_init_ckpt_and_exit`. Family separation should be close to the null distribution. The gap between trained and untrained models quantifies the contribution of pretraining.
-
-Together, these controls support the claim that the family structure of the module
-latent space is produced by pathway-based module design plus pretraining, rather
-than by module design alone or by a generic encoder architecture.
-
-### Control 1: Random Module Assignment (`shuffle_module_genes`)
-
-Add `--shuffle_biomod_genes` to the original command and write to a control output
-directory.
-
-For h5ad data, using `pancreas_scib` as an example:
-
-```bash
-python BioMod/interpretability/biomod_geometry.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl           weight/BioMod.pkl \
-    --biomod_data_dir       model_dependencies/biomod_data_all \
-    --gene_list_path       model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path    model_dependencies/gene2ens.csv \
-    --gene2vec_data_path   model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file        model_dependencies/gene2ens.csv \
-    --data_path            data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --data_mode            h5ad \
-    --gene_name_column     auto \
-    --batch_size           640 \
-    --shuffle_biomod_genes \
-    --output_dir           outputs/interpretability/geometry/pancreas_scib_shuffle
-```
-
-For pretraining npz data:
-
-```bash
-python BioMod/interpretability/biomod_geometry.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl           weight/BioMod.pkl \
-    --biomod_data_dir       model_dependencies/biomod_data_all \
-    --gene_list_path       model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path    model_dependencies/gene2ens.csv \
-    --gene2vec_data_path   model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file        model_dependencies/gene2ens.csv \
-    --data_path            /absolute/path/to/pretraining_cell_npz \
-    --data_mode            pretrain_npz \
-    --max_cells            50000 \
-    --batch_size           640 \
-    --shuffle_biomod_genes \
-    --output_dir           outputs/interpretability/geometry/pretrain_subset_shuffle
-```
-
-How `--shuffle_biomod_genes` works: inside `load_backbone()`, module sizes are
-kept unchanged, all genes are collected into one list, the list is randomly
-shuffled, and genes are reassigned to modules. The seed is controlled by
-`--shuffle_seed`, defaulting to 2021. Model weights are unchanged; only the
-gene-to-module mapping is shuffled.
-
-**Expected result:** kNN purity should drop to about 0.576, the null level, and
-the inter-/intra-family distance ratio should drop to about 1.0. `metrics.json`
-records `"shuffle_module_genes": true`.
-
-### Control 2: Untrained Backbone with Random Initialization
-
-First generate a randomly initialized checkpoint, then run `biomod_geometry.py`
-with it.
-
-A matching randomly initialized checkpoint and configuration are not bundled.
-Create them separately with the pretraining script's `--save_init_ckpt_and_exit`
-option and architecture settings matching the BioMod checkpoint being compared.
-Keep the generated control files in a separate output directory. Then pass their
-paths explicitly; do not use the pretrained BioMod weights for this control.
-
-```bash
-python BioMod/interpretability/biomod_geometry.py \
-    --biomod_checkpoint /absolute/path/to/random_init.pth \
-    --config_pkl /absolute/path/to/random_init.pkl \
-    --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --data_mode h5ad \
-    --gene_name_column auto \
-    --batch_size 8 \
-    --output_dir outputs/interpretability/geometry/pancreas_init
-```
-
-**Expected result:** all family-level metrics should be close to the null
-distribution: kNN purity about 0.576, distance ratio about 1.0, and p-value
-above 0.05. This indicates that a randomly initialized encoder cannot separate
-pathway families.
-
-### Result Comparison Matrix
-
-After running all controls, construct the following table for the paper:
-
-| Condition | kNN purity, cosine | Inter-/intra-family ratio, cosine | Conclusion |
-| --- | --- | --- | --- |
-| **Normal**, trained + pathway module | about 0.72, p<0.005 | 3.8-6.2, p<0.005 | Significant family structure |
-| **Shuffle**, trained + random module | about 0.576 expected | about 1.0 expected | Structure is driven by pathway design |
-| **Init**, untrained + pathway module | about 0.576 expected | about 1.0 expected | Structure is learned during pretraining |
-
-The three experiments jointly show that the family structure of the module latent
-space is generated by both pathway-based module design and pretraining; both are
-necessary.
-
-## Experiment 2: Module Embedding to Pathway Text Alignment (`biomod_text_alignment/`)
-
-**Hypothesis:** the structure of the BioMod module embedding space significantly
-aligns with natural-language pathway semantics. Semantically similar pathways
-are also encoded nearby by the model. Gene-level models cannot reproduce this
-experiment because individual gene names do not carry phrase-level semantics.
-
-**Method:**
-
-1. Parse 541 semantically named modules from `merged_biomods.json`: 50 Hallmark and 491 Reactome.
-2. Encode pathway names with S-PubMedBert-MS-MARCO, a sentence-level fine-tuned PubMedBERT model.
-3. Load `module_embeddings.npy` from Experiment 1, containing cell-averaged module embeddings.
-4. Measure alignment:
-   - **RSA / Mantel test**, the main metric: Spearman correlation between paired distance matrices from the two spaces, with a permutation test.
-   - **CKA**: global representation-structure alignment.
-   - **LOO Ridge retrieval**: cross-modal retrieval recall@k.
-
-**Data requirement:** only `module_embeddings.npy` from Experiment 1 is needed;
-raw expression data are not required.
-
-### Commands
-
-```bash
-# Run from the BioMod repository root.
-
-# Download the sentence encoder.
-# The sentence-transformers loader downloads this model on first use.
-# Network access or an existing local model cache is required.
-
-# Main experiment: RSA + CKA + LOO retrieval.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --n_permutations 1000 \
-    --output_dir outputs/interpretability/text_alignment/pretrain_subset_v2
-
-# Extended analysis: broad Reactome categories and Hallmark deep dive.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_extended.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --descriptions_file BioMod/interpretability/biomod_text_alignment/reactome_descriptions.json \
-    --n_permutations 1000 \
-    --output_dir outputs/interpretability/text_alignment/pretrain_subset_v2_extended
-```
-
-### Output Files
-
-| File | Content |
-| --- | --- |
-| `alignment_metrics.json` | RSA, CKA, retrieval recall@k, and p-values |
-| `rsa_scatter.png` | Paired-distance scatter plot, module space vs text space |
-| `rsa_null_distribution.png` | Mantel test: observed value vs null distribution |
-| `retrieval_ranks.png` | Histogram of retrieval ranks |
-| `retrieval_examples.csv` | Top-5 retrieval results for representative pathways |
-| `extended_metrics.json` | Extended-analysis results: broad-category RSA and Hallmark deep dive |
-| `reactome_category_heatmaps.png` | Distance-matrix comparison at the broad Reactome category level |
-| `hallmarks_deep_dive.png` | Hallmark per-module alignment scores and similarity matrix |
-| `hallmarks_per_module_alignment.csv` | Per-module alignment ranking for 50 Hallmarks |
-
-## Dependencies
-
-`torch`, `numpy`, `pandas`, `scikit-learn`, `matplotlib`, `anndata`, and `scipy`.
-UMAP plots additionally require `umap-learn`; if it is not installed, the script
-warns and skips only UMAP plots. Experiment 2 also requires
-`sentence-transformers`, installable with `pip install sentence-transformers`.

BioMod/interpretability/__init__.py

@@ -1,8 +0,0 @@
-"""Module interpretability experiments for scLong (biomod).
-
-Roadmap (priority order):
-    1. biomod_geometry.py — intrinsic module-embedding geometry (current).
-    2. biomod_text_alignment.py — module ↔ pathway natural-language alignment.
-    3. biomod_importance.py — task-driven ablation + literature enrichment.
-    4. module_vs_gene.py — comparison with gene-level baselines (GenePT, scGPT, ...).
-"""

BioMod/interpretability/attention/README.md

@@ -1,111 +0,0 @@
-# Experiment 6: Module Crosstalk Graph vs Pathway Interaction Database
-
-This experiment tests whether module-module interactions learned by the BioMod
-encoder are consistent with known pathway-pathway relationships.
-
-## Methods
-
-| Method | Principle | Compute cost |
-| --- | --- | --- |
-| A, cosine | Cosine similarity between cell-averaged module embeddings | Low |
-| C, jacobian | Cell-averaged norm of `d encoder_out[j] / d module_emb[i]` | High |
-
-Ground truth: Jaccard index and overlap coefficient of shared genes between
-modules.
-
-## Evaluation Metrics
-
-- Spearman correlation between the model interaction matrix and database overlap matrix.
-- AUROC, using database edges as positives and model scores as predictions.
-- Precision@k, the fraction of top-k model interactions that exist in the database.
-- p-value from a permutation null distribution.
-
-## Run
-
-Run from the BioMod repository root. The h5ad examples use the bundled pancreas
-dataset and BioMod model weights. The optional pretraining `.npz` dataset is not
-included; replace `/absolute/path/to/pretraining_cell_npz` with your own directory.
-
-```bash
-# Run from the BioMod repository root.
-
-# Full run, methods A + C.
-python BioMod/interpretability/attention/biomod_crosstalk.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --data_mode h5ad \
-    --batch_size 640 \
-    --jacobian_cells 200 \
-    --output_dir outputs/interpretability/crosstalk/pancreas_max \
-    --pool_mode max
-
-# Method A only, for quick validation.
-python BioMod/interpretability/attention/biomod_crosstalk.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --data_mode h5ad \
-    --batch_size 64 \
-    --skip_jacobian \
-    --output_dir outputs/interpretability/crosstalk/pancreas_cosine_only
-
-# Pretraining npz data.
-python BioMod/interpretability/attention/biomod_crosstalk.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path /absolute/path/to/pretraining_cell_npz \
-    --data_mode pretrain_npz \
-    --max_cells 5000 \
-    --batch_size 16 \
-    --jacobian_cells 100 \
-    --output_dir outputs/interpretability/crosstalk/pretrain
-```
-
-## Key Arguments
-
-| Argument | Default | Description |
-| --- | --- | --- |
-| `--jacobian_cells` | 200 | Number of cells sampled for Jacobian computation; more is stabler but slower |
-| `--skip_jacobian` | off | Run method A only |
-| `--max_cells` | 5000 | Maximum cell count used by method A |
-| `--null_permutations` | 200 | Number of permutations |
-| `--top_k_fractions` | 0.01,0.05,0.1 | Fractions used for Precision@k |
-
-## Outputs
-
-```text
-output_dir/
-|-- cosine_interaction.npy       # (660, 660) method A interaction matrix
-|-- jacobian_interaction.npy     # (660, 660) method C interaction matrix
-|-- mean_module_embeddings.npy    # (660, 1024) cell-averaged module embeddings
-|-- gt_jaccard.npy               # (660, 660) gene Jaccard overlap
-|-- gt_overlap_coef.npy          # (660, 660) overlap coefficient
-|-- metrics.json                 # All evaluation metrics and p-values
-|-- heatmap_cosine.png           # Method A heatmap, sorted by family
-|-- heatmap_jacobian.png         # Method C heatmap
-|-- heatmap_gt_jaccard.png       # Ground-truth heatmap
-|-- scatter_cosine_vs_jaccard.png
-`-- scatter_jacobian_vs_jaccard.png
-```
-
-## Notes
-
-- Jacobian computation requires 660 backward passes per cell. `--jacobian_cells 200` takes about 10-30 minutes on one A100.
-- Method A has almost no extra overhead, so use `--skip_jacobian` for quick validation first.
-- The current model uses TSSA attention rather than standard softmax attention, so attention weights cannot be read directly. Jacobian is used as a substitute.

BioMod/interpretability/attention/README_gene_attention.md

@@ -1,112 +0,0 @@
-# Gene Attention Interpretability: BioModEmbedding
-
-Analyze gene contributions within modules using the content-based attention
-pooling in BioModEmbedding.
-
-## Analyses
-
-1. Global gene-importance ranking, averaged attention over all cells.
-2. Cell-type specificity with Mann-Whitney U test and fold change.
-3. Marker gene validation with precision@k and recall@k.
-4. Pathway enrichment analysis with GO BP, requiring `gseapy`.
-5. Comparison with DEG results using Spearman correlation.
-
-## Dependencies
-
-```bash
-pip install torch numpy pandas scipy matplotlib seaborn anndata
-# Optional
-pip install gseapy  # Pathway enrichment analysis
-```
-
-## Run Commands
-
-Run from the BioMod repository root. The examples use bundled BioMod model
-weights and pancreas data. Custom DEG and marker CSV files are optional external
-inputs; replace their example paths with your own files before using those flags.
-Set `CELL_TYPE_COLUMN` below to the actual cell-type field in your h5ad metadata;
-the column name has not been assumed for the bundled pancreas file.
-
-### Basic Usage on the Pancreas Dataset with Built-In Markers
-
-```bash
-# Run from the BioMod repository root.
-
-python BioMod/interpretability/attention/biomod_gene_attention.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --celltype_column CELL_TYPE_COLUMN \
-    --batch_size 640 \
-    --top_k_genes 20 \
-    --output_dir outputs/interpretability/gene_attention/pancreas
-```
-
-### With DEG Comparison
-
-```bash
-python BioMod/interpretability/attention/biomod_gene_attention.py \
-    --biomod_checkpoint weight/BioMod.pth \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-    --celltype_column CELL_TYPE_COLUMN \
-    --batch_size 64 \
-    --deg_csv /path/to/deg_results.csv \
-    --output_dir outputs/interpretability/gene_attention/pancreas_with_deg
-```
-
-DEG CSV columns must include `gene`, `cell_type`, `log2fc`, and `pval`.
-
-### Custom Marker Database
-
-Add `--marker_db /path/to/markers.csv` to the complete basic command above and
-change `--output_dir` to `outputs/interpretability/gene_attention/custom_markers`.
-
-Marker CSV columns must include `cell_type` and `gene_symbol`.
-
-### Limit Cell Count for Faster Testing
-
-Add `--max_cells 5000` to the complete basic command above and change
-`--output_dir` to `outputs/interpretability/gene_attention/pancreas_5k`.
-
-## Key Arguments
-
-| Argument | Default | Description |
-| --- | --- | --- |
-| `--batch_size` | 64 | Inference batch size |
-| `--top_k_genes` | 20 | Number of top genes shown for each cell type in visualizations |
-| `--max_cells` | None | Number of cells to subsample |
-| `--marker_db` | None | Custom marker CSV; built-in pancreas markers are used when omitted |
-| `--deg_csv` | None | DEG result CSV; DEG correlation analysis is skipped when omitted |
-| `--celltype_column` | celltype | Cell-type column name in h5ad |
-| `--shuffle_biomod_genes` | off | Randomly shuffle gene-to-module assignment as a negative control |
-
-## Output Files
-
-```text
-output_dir/
-|-- global_gene_importance.csv      # Global gene importance across all modules and genes
-|-- global_top_genes.png            # Top-N gene bar chart
-|-- celltype_<ct>_attention.csv     # Gene attention for each cell type
-|-- celltype_attention_heatmap.png  # Cell type x gene heatmap
-|-- gene_specificity.csv            # Gene specificity with fold change and p-value
-|-- specificity_volcano.png         # Volcano plot
-|-- marker_validation.json          # Marker validation results
-|-- marker_recall.png               # Recall@50/100 bar chart
-|-- deg_correlation.json            # DEG correlation, if --deg_csv is provided
-|-- deg_correlation.png             # Spearman rho bar chart
-|-- enrichment/                     # Pathway enrichment, if gseapy is installed
-|   |-- <ct>_enrichment.csv
-|   `-- enrichment_summary.json
-`-- summary.json                    # Run summary
-```

BioMod/interpretability/attention/biomod_crosstalk.py

@@ -1,606 +0,0 @@
-"""Experiment 6 — Module crosstalk graph vs Reactome/KEGG pathway interactions.
-
-Two complementary methods to quantify module-to-module interaction strength:
-
-  Method A (cosine similarity): Compute cosine similarity between cell-averaged
-      module encoder outputs. Fast, no gradient needed.
-
-  Method C (Jacobian): Compute ||d encoder_out[j] / d module_input[i]|| averaged
-      over cells. Measures causal influence of module i on module j's representation.
-      Requires gradient; more expensive but more principled.
-
-Both interaction matrices are compared against known pathway-pathway relationships
-from Reactome (shared genes / hierarchical overlap) and optionally KEGG.
-
-Metrics:
-  - Spearman correlation between model interaction matrix and DB overlap matrix
-  - Precision@k: among top-k model interactions, fraction present in DB
-  - AUROC: treating DB edges as positives, model scores as predictions
-  - Permutation null p-values for all metrics
-
-Usage (run from the BioMod repository root):
-  python BioMod/interpretability/attention/biomod_crosstalk.py \
-      --biomod_checkpoint weight/BioMod.pth \
-      --config_pkl weight/BioMod.pkl \
-      --biomod_data_dir model_dependencies/biomod_data_all \
-      --gene_list_path model_dependencies/selected_genes_27k.txt \
-      --gene_mapping_path model_dependencies/gene2ens.csv \
-      --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-      --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-      --data_mode h5ad \
-      --batch_size 16 \
-      --jacobian_cells 200 \
-      --output_dir outputs/interpretability/crosstalk/pancreas
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import torch
-from torch.utils.data import DataLoader, Subset
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-from scipy.stats import spearmanr
-from sklearn.metrics import roc_auc_score
-
-HERE = Path(__file__).resolve()
-BIOMOD_SOURCE = HERE.parents[2]
-REPO_ROOT = HERE.parents[3]
-ASSET_ROOT = REPO_ROOT / "model_dependencies"
-for p in (str(BIOMOD_SOURCE), str(REPO_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from cell_annotation.biomod_finetune import (
-    DataPreprocessor,
-    load_backbone,
-    seed_all,
-)
-
-from interpretability.utils import (
-    MatrixRowDataset,
-    NpzCellDataset,
-    build_biomod_metadata,
-    ensure_dir,
-    load_biomod_names,
-)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-logger = logging.getLogger(__name__)
-
-# ---------------------------------------------------------------------------
-# Argparse
-# ---------------------------------------------------------------------------
-
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-
-    p.add_argument("--biomod_checkpoint", "--performer_checkpoint", type=str, required=True, dest='performer_checkpoint')
-    p.add_argument("--config_pkl", type=str, required=True)
-    p.add_argument("--biomod_data_dir", "--module_data_dir", type=str, default=str(ASSET_ROOT / "biomod_data_all"), dest='module_data_dir')
-    p.add_argument("--biomod_json", "--module_json", type=str, default=None, dest='module_json')
-    p.add_argument("--gene_list_path", type=str, default=str(ASSET_ROOT / "selected_genes_27k.txt"))
-    p.add_argument("--gene_mapping_path", type=str, default=str(ASSET_ROOT / "gene2ens.csv"))
-    p.add_argument("--gene2vec_data_path", type=str, default=str(ASSET_ROOT / "selected_gene2vec_27k.npy"))
-    p.add_argument("--gene2ens_file", type=str, default=str(ASSET_ROOT / "gene2ens.csv"))
-    p.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
-    p.add_argument("--shuffle_biomod_genes", "--shuffle_module_genes", action="store_true", dest='shuffle_module_genes')
-    p.add_argument("--shuffle_seed", type=int, default=2021)
-
-    p.add_argument("--data_path", type=str, required=True)
-    p.add_argument("--data_mode", type=str, default="h5ad",
-                   choices=["h5ad", "pretrain_h5ad", "pretrain_npz"])
-    p.add_argument("--cell_num", type=int, default=None)
-    p.add_argument("--gene_name_column", type=str, default="auto")
-    p.add_argument("--dataset_name", type=str, default="crosstalk")
-    p.add_argument("--max_cells", type=int, default=5000,
-                   help="Max cells for method A (cosine). Method C uses --jacobian_cells.")
-
-    p.add_argument("--batch_size", type=int, default=16)
-    p.add_argument("--num_workers", type=int, default=0)
-    p.add_argument("--device", type=str, default="cuda")
-    p.add_argument("--seed", type=int, default=2021)
-    p.add_argument("--pool_mode", type=str, default="mean",
-                   choices=["mean", "max"],
-                   help="Pooling over cells: 'mean' or 'max' (element-wise max).")
-
-    p.add_argument("--jacobian_cells", type=int, default=200,
-                   help="Number of cells for Jacobian computation (expensive).")
-    p.add_argument("--jacobian_batch", type=int, default=4,
-                   help="Cells per Jacobian forward-backward pass.")
-    p.add_argument("--skip_jacobian", action="store_true",
-                   help="Only run method A (cosine), skip method C (Jacobian).")
-
-    p.add_argument("--null_permutations", type=int, default=200)
-    p.add_argument("--top_k_fractions", type=str, default="0.01,0.05,0.1",
-                   help="Comma-separated fractions for precision@k.")
-
-    p.add_argument("--output_dir", type=str, required=True)
-    return p.parse_args()
-
-
-# ---------------------------------------------------------------------------
-# Data loading (reuses patterns from biomod_geometry.py)
-# ---------------------------------------------------------------------------
-
-def build_dataset(args, preprocessor):
-    """Build dataset and optional reindex_state."""
-    if args.data_mode == "h5ad":
-        import anndata as ad
-        adata, reindex_state = preprocessor.preprocess_adata(
-            adata_path=args.data_path,
-            min_genes=0, min_cells=0,
-            max_cells=args.max_cells,
-            dataset_name=args.dataset_name,
-            gene_name_column=args.gene_name_column,
-        )
-        dataset = MatrixRowDataset(adata.X)
-        n_cells = adata.n_obs
-    elif args.data_mode == "pretrain_h5ad":
-        import anndata as ad
-        adata = ad.read_h5ad(args.data_path)
-        if args.max_cells and adata.n_obs > args.max_cells:
-            rng = np.random.default_rng(args.seed)
-            chosen = np.sort(rng.choice(adata.n_obs, size=args.max_cells, replace=False))
-            adata = adata[chosen].copy()
-        dataset = MatrixRowDataset(adata.X)
-        n_cells = adata.n_obs
-        reindex_state = None
-    elif args.data_mode == "pretrain_npz":
-        dataset = NpzCellDataset(
-            data_dir=args.data_path,
-            cell_num=args.cell_num,
-            max_cells=args.max_cells,
-            seed=args.seed,
-        )
-        n_cells = len(dataset)
-        reindex_state = None
-    else:
-        raise ValueError(f"Unsupported data_mode: {args.data_mode}")
-    return dataset, reindex_state, n_cells
-
-
-# ---------------------------------------------------------------------------
-# Method A: Cosine similarity of cell-averaged module embeddings
-# ---------------------------------------------------------------------------
-
-@torch.no_grad()
-def compute_cosine_interaction(backbone, loader, preprocessor, reindex_state,
-                               device, num_modules, base_dim, expected_genes,
-                               pool_mode="mean"):
-    """Compute pooled module embeddings, return cosine similarity matrix."""
-    backbone.eval()
-    if pool_mode == "mean":
-        agg_emb = torch.zeros(num_modules, base_dim, device=device, dtype=torch.float64)
-    else:
-        agg_emb = torch.full((num_modules, base_dim), -float("inf"), device=device, dtype=torch.float64)
-    n_cells = 0
-
-    for batch_idx, x in enumerate(loader):
-        x = x.to(device)
-        if reindex_state is not None:
-            x = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
-        out = backbone(x, mask_ratio=0.0, return_encodings=True)
-        enc = out["module_encodings"].to(torch.float64)  # (B, num_modules, D)
-        if pool_mode == "mean":
-            agg_emb += enc.sum(dim=0)
-        else:
-            batch_max = enc.max(dim=0).values  # (num_modules, D)
-            agg_emb = torch.maximum(agg_emb, batch_max)
-        n_cells += enc.shape[0]
-        if (batch_idx + 1) % 50 == 0:
-            logger.info("  [cosine/%s] batch %d, %d cells", pool_mode, batch_idx + 1, n_cells)
-
-    if pool_mode == "mean":
-        pooled_emb = (agg_emb / max(n_cells, 1)).float().cpu().numpy()
-    else:
-        pooled_emb = agg_emb.float().cpu().numpy()
-
-    norms = np.linalg.norm(pooled_emb, axis=1, keepdims=True)
-    norms = np.clip(norms, 1e-8, None)
-    normed = pooled_emb / norms
-    cos_sim = normed @ normed.T
-    np.fill_diagonal(cos_sim, 0.0)
-    logger.info("Method A (%s) done: %d cells, interaction matrix shape %s", pool_mode, n_cells, cos_sim.shape)
-    return cos_sim, pooled_emb
-
-
-# ---------------------------------------------------------------------------
-# Method C: Jacobian-based module interaction
-# ---------------------------------------------------------------------------
-
-def compute_jacobian_interaction(backbone, loader, preprocessor, reindex_state,
-                                 device, num_modules, base_dim, expected_genes,
-                                 max_cells, batch_size):
-    """Compute average ||d encoder_out[j] / d module_emb[i]|| interaction matrix."""
-    backbone.eval()
-    for param in backbone.parameters():
-        param.requires_grad_(False)
-
-    interaction = torch.zeros(num_modules, num_modules, device=device, dtype=torch.float64)
-    n_cells = 0
-
-    for x in loader:
-        if n_cells >= max_cells:
-            break
-        x = x.to(device)
-        if reindex_state is not None:
-            x = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
-
-        actual_batch = min(x.shape[0], max_cells - n_cells)
-        x = x[:actual_batch]
-
-        for cell_idx in range(x.shape[0]):
-            cell = x[cell_idx:cell_idx+1]
-            J = _compute_single_cell_jacobian(backbone, cell, num_modules, base_dim, device)
-            interaction += J
-            n_cells += 1
-            if n_cells % 20 == 0:
-                logger.info("  [jacobian] %d / %d cells", n_cells, max_cells)
-            if n_cells >= max_cells:
-                break
-
-    interaction /= max(n_cells, 1)
-    result = interaction.cpu().numpy()
-    np.fill_diagonal(result, 0.0)
-    logger.info("Method C done: %d cells, interaction matrix shape %s", n_cells, result.shape)
-    return result
-
-
-def _compute_single_cell_jacobian(backbone, cell, num_modules, base_dim, device):
-    """For one cell, compute (num_modules, num_modules) Jacobian norm matrix.
-
-    J[i, j] = ||d encoder_out[j] / d module_emb[i]||_F
-
-    Strategy: run module_embed to get intermediate, detach, enable grad,
-    then replicate the encoder forward path and backprop from each output module.
-    """
-    J = torch.zeros(num_modules, num_modules, device=device, dtype=torch.float32)
-
-    with torch.no_grad():
-        x_raw, _ = backbone.module_embed(cell)  # (1, num_modules, D)
-
-    x_modules = x_raw.detach().clone().requires_grad_(True)
-
-    with torch.enable_grad():
-        if backbone.mask_style == "replace":
-            x = backbone._add_biomod_position_embedding(x_modules)
-        else:
-            x = backbone._add_biomod_position_embedding(x_modules)
-
-        if backbone.use_cls and backbone.cls_token is not None:
-            cls_token = backbone.cls_token + backbone.pos_emb[:, :1, :]
-            cls_tokens = cls_token.expand(x.shape[0], -1, -1)
-            x = torch.cat((cls_tokens, x), dim=1)
-
-        x_emb = backbone.dropout(x)
-        layer_pos_emb = backbone.layer_pos_emb(x_emb)
-        enc_out = backbone.encoder(x_emb, pos_emb=layer_pos_emb)
-        enc_out = backbone.encoder_norm(enc_out)  # (1, L, D)
-
-        if backbone.use_cls and backbone.cls_token is not None:
-            enc_out = enc_out[:, 1:, :]  # skip CLS
-
-        for j in range(num_modules):
-            if x_modules.grad is not None:
-                x_modules.grad.zero_()
-            scalar = enc_out[0, j].pow(2).sum()
-            scalar.backward(retain_graph=(j < num_modules - 1))
-            grad = x_modules.grad[0]  # (num_modules, D)
-            J[:, j] = grad.norm(dim=1)
-
-    return J.to(torch.float64)
-
-
-# ---------------------------------------------------------------------------
-# Ground truth: pathway overlap from module gene sets
-# ---------------------------------------------------------------------------
-
-def compute_pathway_overlap(module_names: List[str], module_gene_dict: dict,
-                            metric: str = "jaccard") -> np.ndarray:
-    """Compute pairwise pathway overlap matrix from shared genes.
-
-    metric: 'jaccard' (|A∩B| / |A∪B|) or 'overlap_coef' (|A∩B| / min(|A|,|B|))
-    """
-    n = len(module_names)
-    gene_sets = []
-    for name in module_names:
-        genes = module_gene_dict.get(name, [])
-        gene_sets.append(set(genes))
-
-    overlap = np.zeros((n, n), dtype=np.float32)
-    for i in range(n):
-        for j in range(i + 1, n):
-            inter = len(gene_sets[i] & gene_sets[j])
-            if inter == 0:
-                continue
-            if metric == "jaccard":
-                union = len(gene_sets[i] | gene_sets[j])
-                score = inter / union if union > 0 else 0.0
-            else:
-                min_size = min(len(gene_sets[i]), len(gene_sets[j]))
-                score = inter / min_size if min_size > 0 else 0.0
-            overlap[i, j] = score
-            overlap[j, i] = score
-    return overlap
-
-
-# ---------------------------------------------------------------------------
-# Evaluation metrics
-# ---------------------------------------------------------------------------
-
-def evaluate_interaction(model_matrix: np.ndarray, gt_matrix: np.ndarray,
-                         top_k_fracs: List[float], n_permutations: int = 200):
-    """Compare model interaction matrix against ground truth overlap.
-
-    Returns dict with spearman_r, auroc, precision@k, and permutation p-values.
-    """
-    n = model_matrix.shape[0]
-    triu_idx = np.triu_indices(n, k=1)
-    model_flat = model_matrix[triu_idx]
-    gt_flat = gt_matrix[triu_idx]
-
-    has_positive = gt_flat.max() > 0
-    has_variation = model_flat.std() > 1e-10
-
-    rho, _ = spearmanr(model_flat, gt_flat) if has_variation else (0.0, 1.0)
-
-    gt_binary = (gt_flat > 0).astype(np.float32)
-    n_pos = gt_binary.sum()
-    if n_pos > 0 and n_pos < len(gt_binary) and has_variation:
-        auroc = roc_auc_score(gt_binary, model_flat)
-    else:
-        auroc = 0.5
-
-    precisions = {}
-    n_pairs = len(model_flat)
-    for frac in top_k_fracs:
-        k = max(1, int(n_pairs * frac))
-        top_k_idx = np.argsort(model_flat)[-k:]
-        precisions[f"precision@{frac}"] = float(gt_binary[top_k_idx].mean())
-
-    null_rhos = []
-    null_aurocs = []
-    rng = np.random.default_rng(42)
-    for _ in range(n_permutations):
-        perm = rng.permutation(n)
-        perm_matrix = model_matrix[perm][:, perm]
-        perm_flat = perm_matrix[triu_idx]
-        r, _ = spearmanr(perm_flat, gt_flat) if has_variation else (0.0, 1.0)
-        null_rhos.append(r)
-        if n_pos > 0 and n_pos < len(gt_binary):
-            null_aurocs.append(roc_auc_score(gt_binary, perm_flat))
-        else:
-            null_aurocs.append(0.5)
-
-    p_rho = (np.sum(np.array(null_rhos) >= rho) + 1) / (n_permutations + 1)
-    p_auroc = (np.sum(np.array(null_aurocs) >= auroc) + 1) / (n_permutations + 1)
-
-    return {
-        "spearman_r": float(rho),
-        "spearman_p_perm": float(p_rho),
-        "auroc": float(auroc),
-        "auroc_p_perm": float(p_auroc),
-        "null_rho_mean": float(np.mean(null_rhos)),
-        "null_rho_std": float(np.std(null_rhos)),
-        **precisions,
-    }
-
-
-# ---------------------------------------------------------------------------
-# Visualization
-# ---------------------------------------------------------------------------
-
-def plot_interaction_matrix(matrix: np.ndarray, title: str, path: Path,
-                            module_families: Optional[List[str]] = None):
-    """Plot heatmap of interaction matrix, optionally sorted by family."""
-    if module_families is not None:
-        order = np.argsort(module_families)
-        matrix = matrix[order][:, order]
-        sorted_families = [module_families[i] for i in order]
-    else:
-        sorted_families = None
-
-    fig, ax = plt.subplots(figsize=(10, 9))
-    vmax = np.percentile(np.abs(matrix), 99)
-    im = ax.imshow(matrix, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="equal")
-    ax.set_title(title, fontsize=12)
-    ax.set_xlabel("Module j")
-    ax.set_ylabel("Module i")
-    plt.colorbar(im, ax=ax, fraction=0.046)
-
-    if sorted_families is not None:
-        unique_fams = []
-        boundaries = [0]
-        prev = sorted_families[0]
-        for idx, f in enumerate(sorted_families):
-            if f != prev:
-                unique_fams.append(prev)
-                boundaries.append(idx)
-                prev = f
-        unique_fams.append(prev)
-        boundaries.append(len(sorted_families))
-        for b in boundaries[1:-1]:
-            ax.axhline(b - 0.5, color="black", linewidth=0.3, alpha=0.5)
-            ax.axvline(b - 0.5, color="black", linewidth=0.3, alpha=0.5)
-
-    plt.tight_layout()
-    fig.savefig(path, dpi=150, bbox_inches="tight")
-    plt.close(fig)
-
-
-def plot_scatter(model_flat, gt_flat, title, path):
-    """Scatter plot of model interaction vs ground truth overlap."""
-    fig, ax = plt.subplots(figsize=(6, 5))
-    subsample = min(len(model_flat), 50000)
-    if subsample < len(model_flat):
-        rng = np.random.default_rng(0)
-        idx = rng.choice(len(model_flat), size=subsample, replace=False)
-        model_flat = model_flat[idx]
-        gt_flat = gt_flat[idx]
-    ax.scatter(gt_flat, model_flat, alpha=0.1, s=2, rasterized=True)
-    ax.set_xlabel("Ground truth overlap (Jaccard)")
-    ax.set_ylabel("Model interaction score")
-    ax.set_title(title)
-    plt.tight_layout()
-    fig.savefig(path, dpi=150, bbox_inches="tight")
-    plt.close(fig)
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-def main():
-    args = parse_args()
-    seed_all(args.seed)
-    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-    output_dir = Path(args.output_dir)
-    ensure_dir(output_dir)
-
-    # Resolve all paths to absolute (relative to CWD) so that load_backbone's
-    # internal resolve_path (which resolves relative to cell_annotation/) is bypassed.
-    for attr in ("performer_checkpoint", "config_pkl", "module_data_dir",
-                 "gene_list_path", "gene_mapping_path", "gene2vec_data_path",
-                 "gene2ens_file", "data_path"):
-        val = getattr(args, attr, None)
-        if val is not None and not Path(val).is_absolute():
-            setattr(args, attr, str(Path(val).resolve()))
-
-    logger.info("Loading backbone...")
-    backbone, biomod_data, model_config = load_backbone(args, device)
-    backbone.eval()
-
-    num_modules = len(biomod_data)
-    base_dim = int(model_config["base_dim"])
-    expected_genes = int(model_config.get("num_genes", 27874))
-    logger.info("Model: %d modules, dim=%d, genes=%d", num_modules, base_dim, expected_genes)
-
-    module_json_path = args.module_json or str(Path(args.module_data_dir) / "merged_biomods.json")
-    module_names = load_biomod_names(module_json_path)
-    with open(module_json_path, "r", encoding="utf-8") as f:
-        module_gene_dict = json.load(f)
-
-    preprocessor = None
-    if args.data_mode == "h5ad":
-        preprocessor = DataPreprocessor(
-            gene2ens_file=args.gene2ens_file,
-            gene2ens_column=args.gene2ens_column,
-            model_genes_file=args.gene_list_path,
-        )
-        preprocessor.load_gene_mapping()
-        preprocessor.load_model_genes()
-
-    dataset, reindex_state, n_cells = build_dataset(args, preprocessor)
-    logger.info("Dataset: %d cells", n_cells)
-
-    # --- Method A: Cosine ---
-    logger.info("=" * 60)
-    logger.info("Method A: Cosine similarity of module embeddings")
-    logger.info("=" * 60)
-    loader_a = DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
-                          num_workers=args.num_workers, pin_memory=True)
-    cos_matrix, mean_emb = compute_cosine_interaction(
-        backbone, loader_a, preprocessor, reindex_state,
-        device, num_modules, base_dim, expected_genes,
-        pool_mode=args.pool_mode,
-    )
-
-    # --- Method C: Jacobian ---
-    jac_matrix = None
-    if not args.skip_jacobian:
-        logger.info("=" * 60)
-        logger.info("Method C: Jacobian-based module interaction")
-        logger.info("=" * 60)
-        jac_n = min(args.jacobian_cells, n_cells)
-        if jac_n < n_cells:
-            rng = np.random.default_rng(args.seed)
-            jac_indices = np.sort(rng.choice(n_cells, size=jac_n, replace=False))
-            jac_dataset = Subset(dataset, jac_indices.tolist())
-        else:
-            jac_dataset = dataset
-        loader_c = DataLoader(jac_dataset, batch_size=args.jacobian_batch, shuffle=False,
-                              num_workers=0, pin_memory=True)
-        jac_matrix = compute_jacobian_interaction(
-            backbone, loader_c, preprocessor, reindex_state,
-            device, num_modules, base_dim, expected_genes,
-            max_cells=jac_n, batch_size=args.jacobian_batch,
-        )
-
-    # --- Ground truth overlap ---
-    logger.info("Computing pathway gene overlap (Jaccard)...")
-    gt_jaccard = compute_pathway_overlap(module_names, module_gene_dict, metric="jaccard")
-    gt_overlap = compute_pathway_overlap(module_names, module_gene_dict, metric="overlap_coef")
-
-    # --- Evaluate ---
-    top_k_fracs = [float(f) for f in args.top_k_fractions.split(",")]
-    results = {}
-
-    logger.info("Evaluating Method A vs Jaccard overlap...")
-    results["cosine_vs_jaccard"] = evaluate_interaction(
-        cos_matrix, gt_jaccard, top_k_fracs, args.null_permutations)
-    results["cosine_vs_overlap_coef"] = evaluate_interaction(
-        cos_matrix, gt_overlap, top_k_fracs, args.null_permutations)
-
-    if jac_matrix is not None:
-        logger.info("Evaluating Method C vs Jaccard overlap...")
-        jac_sym = (jac_matrix + jac_matrix.T) / 2
-        results["jacobian_vs_jaccard"] = evaluate_interaction(
-            jac_sym, gt_jaccard, top_k_fracs, args.null_permutations)
-        results["jacobian_vs_overlap_coef"] = evaluate_interaction(
-            jac_sym, gt_overlap, top_k_fracs, args.null_permutations)
-
-    # --- Save ---
-    np.save(output_dir / "cosine_interaction.npy", cos_matrix)
-    np.save(output_dir / "mean_module_embeddings.npy", mean_emb)
-    np.save(output_dir / "gt_jaccard.npy", gt_jaccard)
-    np.save(output_dir / "gt_overlap_coef.npy", gt_overlap)
-    if jac_matrix is not None:
-        np.save(output_dir / "jacobian_interaction.npy", jac_matrix)
-
-    with open(output_dir / "metrics.json", "w", encoding="utf-8") as f:
-        json.dump(results, f, indent=2, ensure_ascii=False)
-
-    # --- Plots ---
-    meta = build_biomod_metadata(module_names)
-    families = [row["family"] for row in meta]
-
-    plot_interaction_matrix(cos_matrix, "Method A: Cosine Similarity",
-                           output_dir / "heatmap_cosine.png", families)
-    if jac_matrix is not None:
-        plot_interaction_matrix(jac_sym, "Method C: Jacobian Interaction",
-                               output_dir / "heatmap_jacobian.png", families)
-    plot_interaction_matrix(gt_jaccard, "Ground Truth: Jaccard Gene Overlap",
-                           output_dir / "heatmap_gt_jaccard.png", families)
-
-    triu_idx = np.triu_indices(num_modules, k=1)
-    plot_scatter(cos_matrix[triu_idx], gt_jaccard[triu_idx],
-                 "Cosine vs Jaccard", output_dir / "scatter_cosine_vs_jaccard.png")
-    if jac_matrix is not None:
-        plot_scatter(jac_sym[triu_idx], gt_jaccard[triu_idx],
-                     "Jacobian vs Jaccard", output_dir / "scatter_jacobian_vs_jaccard.png")
-
-    # --- Print summary ---
-    logger.info("=" * 60)
-    logger.info("Results Summary")
-    logger.info("=" * 60)
-    for method_key, metrics in results.items():
-        logger.info("  %s:", method_key)
-        logger.info("    Spearman r = %.4f (p_perm = %.4f)", metrics["spearman_r"], metrics["spearman_p_perm"])
-        logger.info("    AUROC = %.4f (p_perm = %.4f)", metrics["auroc"], metrics["auroc_p_perm"])
-        for k in top_k_fracs:
-            logger.info("    Precision@%.2f = %.4f", k, metrics[f"precision@{k}"])
-    logger.info("Output saved to: %s", output_dir)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/attention/biomod_gene_attention.py

@@ -1,1093 +0,0 @@
-"""BioModEmbedding Attention-based Gene Interpretability Analysis (Streaming).
-
-Interpretability analysis using single-head content-based attention pooling in BioModEmbedding.
-Streaming computing mode: memory usage O(n_celltypes × n_genes).
-
-Analysis content:
-  1. Global gene importance ranking
-  2. Cell type specificity (z-score, Cohen's d, Welch's t)
-  3. Marker gene verification
-  4. Attention-Expression correlation
-  5. Random baseline comparison (KL divergence)
-  6. Attention concentration (Gini coefficient)
-  7. Pathway Enrichment Analysis
-  8. DEG comparison
-
-Usage (run from the BioMod repository root):
-  python BioMod/interpretability/attention/biomod_gene_attention.py \
-      --biomod_checkpoint weight/BioMod.pth \
-      --config_pkl weight/BioMod.pkl \
-      --biomod_data_dir model_dependencies/biomod_data_all \
-      --gene_list_path model_dependencies/selected_genes_27k.txt \
-      --gene_mapping_path model_dependencies/gene2ens.csv \
-      --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-      --gene2ens_file model_dependencies/gene2ens.csv \
-      --data_path data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad \
-      --celltype_column CELL_TYPE_COLUMN \
-      --batch_size 64 \
-      --output_dir outputs/interpretability/gene_attention/pancreas
-
-Replace CELL_TYPE_COLUMN with the cell-type field in the input h5ad metadata.
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-import torch
-from torch.utils.data import DataLoader
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-import seaborn as sns
-
-from scipy.stats import spearmanr, t as t_dist
-
-HERE = Path(__file__).resolve()
-BIOMOD_SOURCE = HERE.parents[2]
-REPO_ROOT = HERE.parents[3]
-ASSET_ROOT = REPO_ROOT / "model_dependencies"
-for p in (str(BIOMOD_SOURCE), str(REPO_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-# Publication utilities
-BIOMOD_IMPORTANCE_DIR = BIOMOD_SOURCE / "interpretability" / "biomod_importance"
-if str(BIOMOD_IMPORTANCE_DIR) not in sys.path:
-    sys.path.insert(0, str(BIOMOD_IMPORTANCE_DIR))
-
-from pub_utils import setup_publication_style, savefig as pub_savefig
-
-from cell_annotation.biomod_finetune import (
-    DataPreprocessor,
-    load_backbone,
-    seed_all,
-)
-from interpretability.utils import (
-    MatrixRowDataset,
-    build_biomod_metadata,
-    ensure_dir,
-    load_biomod_names,
-)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# Argparse
-# ---------------------------------------------------------------------------
-
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-
-    p.add_argument("--biomod_checkpoint", "--performer_checkpoint", type=str, required=True, dest='performer_checkpoint')
-    p.add_argument("--config_pkl", type=str, required=True)
-    p.add_argument("--biomod_data_dir", "--module_data_dir", type=str, default=str(ASSET_ROOT / "biomod_data_all"), dest='module_data_dir')
-    p.add_argument("--biomod_json", "--module_json", type=str, default=None, dest='module_json')
-    p.add_argument("--gene_list_path", type=str, default=str(ASSET_ROOT / "selected_genes_27k.txt"))
-    p.add_argument("--gene_mapping_path", type=str, default=str(ASSET_ROOT / "gene2ens.csv"))
-    p.add_argument("--gene2vec_data_path", type=str, default=str(ASSET_ROOT / "selected_gene2vec_27k.npy"))
-    p.add_argument("--gene2ens_file", type=str, default=str(ASSET_ROOT / "gene2ens.csv"))
-    p.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
-    p.add_argument("--shuffle_biomod_genes", "--shuffle_module_genes", action="store_true", dest='shuffle_module_genes')
-    p.add_argument("--shuffle_seed", type=int, default=2021)
-
-    p.add_argument("--data_path", type=str, required=True)
-    p.add_argument("--celltype_column", type=str, default="celltype")
-    p.add_argument("--gene_name_column", type=str, default="auto")
-    p.add_argument("--dataset_name", type=str, default="gene_attention")
-    p.add_argument("--max_cells", type=int, default=None)
-
-    p.add_argument("--batch_size", type=int, default=64)
-    p.add_argument("--num_workers", type=int, default=0)
-    p.add_argument("--device", type=str, default="cuda")
-    p.add_argument("--seed", type=int, default=2021)
-
-    p.add_argument("--top_k_genes", type=int, default=20)
-    p.add_argument("--marker_db", type=str, default=None)
-    p.add_argument("--deg_csv", type=str, default=None)
-    p.add_argument("--min_cells", type=int, default=30)
-
-    p.add_argument("--output_dir", type=str, required=True)
-    return p.parse_args()
-
-
-# ---------------------------------------------------------------------------
-# Built-in pancreas marker genes (from PanglaoDB + CellMarker)
-# ---------------------------------------------------------------------------
-
-PANCREAS_MARKERS: Dict[str, List[str]] = {
-    "alpha": ["GCG", "ARX", "IRX2", "TTR", "LOXL4", "GC", "CRYBA2", "DPP4", "IRX1", "PCSK2"],
-    "beta": ["INS", "IAPP", "MAFA", "NKX6-1", "PDX1", "HADH", "SLC30A8", "NPTX2", "DLK1", "PCSK1"],
-    "delta": ["SST", "HHEX", "RBP4", "LEPR", "PCSK1", "SEC11C"],
-    "gamma": ["PPY", "SERTM1", "CARTPT", "MEIS2", "ETV1", "ABCC9"],
-    "ductal": ["KRT19", "KRT7", "SPP1", "CFTR", "MUC1", "AMBP", "ANXA4", "BICC1", "SOX9"],
-    "acinar": ["PRSS1", "CELA3A", "CELA2A", "CPA1", "CPA2", "PNLIP", "CEL", "CLPS", "REG1A", "CTRB1"],
-    "activated_stellate": ["COL1A1", "COL1A2", "COL3A1", "ACTA2", "TAGLN", "RGS5", "PDGFRB"],
-    "endothelial": ["PECAM1", "VWF", "CDH5", "KDR", "ENG", "PLVAP", "ESAM"],
-    "macrophage": ["CD68", "CD163", "CSF1R", "MARCO", "MSR1", "C1QA", "C1QB"],
-    "epsilon": ["GHRL", "PCSK1", "ARX"],
-}
-
-PANCREAS_COLORS: Dict[str, str] = {
-    "alpha": "#E41A1C", "beta": "#377EB8", "delta": "#4DAF4A",
-    "gamma": "#984EA3", "ductal": "#FF7F00", "acinar": "#A65628",
-    "activated_stellate": "#F781BF", "endothelial": "#66C2A5",
-    "macrophage": "#FC8D62", "epsilon": "#8DA0CB", "t_cell": "#E78AC3",
-    "schwann": "#B3B3B3", "mast": "#FDBF6F", "quiescent_stellate": "#CAB2D6",
-}
-
-
-# ---------------------------------------------------------------------------
-# Streaming attention extraction + statistics accumulation
-# ---------------------------------------------------------------------------
-
-@torch.no_grad()
-def extract_attention_streaming(
-    backbone, loader, preprocessor, reindex_state, device,
-    cell_labels: np.ndarray, num_modules: int, module_sizes: List[int],
-) -> Tuple:
-    """Extract attention and accumulate statistics in streaming fashion.
-
-    Returns:
-        global_mean, global_std, ct_mean, ct_std, ct_count, n_cells,
-        expr_mean, attn_expr_prod_sum
-    """
-    backbone.eval()
-    module_embed = backbone.module_embed
-
-    global_sum = {i: np.zeros(module_sizes[i], dtype=np.float64) for i in range(num_modules)}
-    global_sum_sq = {i: np.zeros(module_sizes[i], dtype=np.float64) for i in range(num_modules)}
-    expr_sum = {i: np.zeros(module_sizes[i], dtype=np.float64) for i in range(num_modules)}
-    attn_expr_prod_sum = {i: np.zeros(module_sizes[i], dtype=np.float64) for i in range(num_modules)}
-
-    celltypes = np.unique(cell_labels)
-    ct_sum = {ct: {i: np.zeros(module_sizes[i], dtype=np.float64) for i in range(num_modules)} for ct in celltypes}
-    ct_sum_sq = {ct: {i: np.zeros(module_sizes[i], dtype=np.float64) for i in range(num_modules)} for ct in celltypes}
-    ct_count = {ct: 0 for ct in celltypes}
-
-    n_cells = 0
-    cell_idx = 0
-
-    for batch_idx, x in enumerate(loader):
-        x = x.to(device)
-        if reindex_state is not None:
-            x = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
-
-        batch_size = x.shape[0]
-        batch_labels = cell_labels[cell_idx:cell_idx + batch_size]
-        cell_idx += batch_size
-
-        for module_i, gene_idx in enumerate(module_embed.module_sets):
-            gene_idx_dev = gene_idx.to(device)
-            expr_slice = torch.index_select(x, dim=1, index=gene_idx_dev)
-
-            gene_emb = module_embed.gene_embed(gene_idx_dev)
-            gene_emb = gene_emb.unsqueeze(0).expand(batch_size, -1, -1)
-            gene_tok = module_embed._fuse_expr_and_gene_emb(expr_slice, gene_emb)
-            gene_tok = module_embed.proj(gene_tok)
-
-            attn = module_embed.attn_score(gene_tok).squeeze(-1)
-            attn = torch.softmax(attn, dim=1)
-            attn_np = attn.cpu().numpy().astype(np.float64)
-
-            # Expression for correlation
-            expr_np = expr_slice.cpu().numpy().astype(np.float64)
-            expr_sum[module_i] += expr_np.sum(axis=0)
-            attn_expr_prod_sum[module_i] += (attn_np * expr_np).sum(axis=0)
-
-            # Global accumulation
-            global_sum[module_i] += attn_np.sum(axis=0)
-            global_sum_sq[module_i] += (attn_np ** 2).sum(axis=0)
-
-            # Per-celltype accumulation
-            for ct in celltypes:
-                mask = batch_labels == ct
-                if mask.any():
-                    ct_attn = attn_np[mask]
-                    ct_sum[ct][module_i] += ct_attn.sum(axis=0)
-                    ct_sum_sq[ct][module_i] += (ct_attn ** 2).sum(axis=0)
-
-        for ct in celltypes:
-            ct_count[ct] += int((batch_labels == ct).sum())
-
-        n_cells += batch_size
-        if (batch_idx + 1) % 50 == 0:
-            logger.info("  [attention] batch %d, %d cells processed", batch_idx + 1, n_cells)
-
-    # Compute means and stds
-    global_mean = {}
-    global_std = {}
-    for i in range(num_modules):
-        global_mean[i] = (global_sum[i] / n_cells).astype(np.float32)
-        variance = global_sum_sq[i] / n_cells - global_mean[i].astype(np.float64) ** 2
-        global_std[i] = np.sqrt(np.maximum(variance, 0)).astype(np.float32)
-
-    ct_mean = {}
-    ct_std = {}
-    for ct in celltypes:
-        ct_mean[ct] = {}
-        ct_std[ct] = {}
-        for i in range(num_modules):
-            if ct_count[ct] > 0:
-                ct_mean[ct][i] = (ct_sum[ct][i] / ct_count[ct]).astype(np.float32)
-            else:
-                ct_mean[ct][i] = np.zeros(module_sizes[i], dtype=np.float32)
-            if ct_count[ct] > 1:
-                var = ct_sum_sq[ct][i] / ct_count[ct] - ct_mean[ct][i].astype(np.float64) ** 2
-                ct_std[ct][i] = np.sqrt(np.maximum(var, 0)).astype(np.float32)
-            else:
-                ct_std[ct][i] = np.zeros(module_sizes[i], dtype=np.float32)
-
-    expr_mean = {i: (expr_sum[i] / n_cells).astype(np.float32) for i in range(num_modules)}
-
-    logger.info("Extracted attention for %d cells across %d modules (streaming)", n_cells, num_modules)
-    return global_mean, global_std, ct_mean, ct_std, ct_count, n_cells, expr_mean, attn_expr_prod_sum
-
-
-# ---------------------------------------------------------------------------
-# Analysis functions
-# ---------------------------------------------------------------------------
-
-def build_global_importance_df(
-    global_mean: Dict[int, np.ndarray],
-    global_std: Dict[int, np.ndarray],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-) -> pd.DataFrame:
-    """Build global gene importance DataFrame from streaming means."""
-    rows = []
-    for module_i in range(num_modules):
-        mean_attn = global_mean[module_i]
-        std_attn = global_std[module_i]
-        genes = module_gene_names[module_i]
-        rank_order = np.argsort(-mean_attn)
-        for g_idx, gene in enumerate(genes):
-            rows.append({
-                "module_idx": module_i,
-                "gene": gene,
-                "mean_attention": float(mean_attn[g_idx]),
-                "std_attention": float(std_attn[g_idx]),
-                "rank_in_module": int(np.where(rank_order == g_idx)[0][0]) + 1,
-                "module_size": len(genes),
-            })
-    df = pd.DataFrame(rows)
-    df["normalized_rank"] = df["rank_in_module"] / df["module_size"]
-    return df
-
-
-def build_celltype_dfs(
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    ct_count: Dict[str, int],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-) -> Dict[str, pd.DataFrame]:
-    """Build per-celltype DataFrames from streaming means."""
-    ct_dfs = {}
-    for ct, module_means in ct_mean.items():
-        rows = []
-        for module_i in range(num_modules):
-            mean_attn = module_means[module_i]
-            genes = module_gene_names[module_i]
-            for g_idx, gene in enumerate(genes):
-                rows.append({
-                    "module_idx": module_i,
-                    "gene": gene,
-                    "mean_attention": float(mean_attn[g_idx]),
-                    "n_cells": ct_count[ct],
-                })
-        ct_dfs[ct] = pd.DataFrame(rows)
-    return ct_dfs
-
-
-def compute_specificity_from_means(
-    global_mean: Dict[int, np.ndarray],
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-) -> pd.DataFrame:
-    """Fold-change specificity (legacy): max_ct(mean_attn_ct) / global_mean_attn."""
-    celltypes = list(ct_mean.keys())
-    rows = []
-    for module_i in range(num_modules):
-        genes = module_gene_names[module_i]
-        g_mean = global_mean[module_i]
-        for g_idx, gene in enumerate(genes):
-            ct_scores = {ct: float(ct_mean[ct][module_i][g_idx]) for ct in celltypes}
-            top_ct = max(ct_scores, key=ct_scores.get)
-            top_score = ct_scores[top_ct]
-            fold_change = top_score / max(float(g_mean[g_idx]), 1e-10)
-            rows.append({
-                "module_idx": module_i, "gene": gene, "top_celltype": top_ct,
-                "top_ct_mean_attn": top_score, "global_mean_attn": float(g_mean[g_idx]),
-                "fold_change": fold_change,
-            })
-    return pd.DataFrame(rows)
-
-
-def compute_specificity_zscore(
-    global_mean: Dict[int, np.ndarray],
-    global_std: Dict[int, np.ndarray],
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    ct_count: Dict[str, int],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-    min_cells: int = 30,
-) -> pd.DataFrame:
-    """Z-score specificity: (ct_mean - global_mean) / global_std."""
-    valid_cts = [ct for ct, n in ct_count.items() if n >= min_cells]
-    rows = []
-    for module_i in range(num_modules):
-        genes = module_gene_names[module_i]
-        g_mean = global_mean[module_i]
-        g_std = global_std[module_i]
-        for g_idx, gene in enumerate(genes):
-            std_val = max(float(g_std[g_idx]), 1e-8)
-            ct_zscores = {}
-            for ct in valid_cts:
-                z = (float(ct_mean[ct][module_i][g_idx]) - float(g_mean[g_idx])) / std_val
-                ct_zscores[ct] = z
-            top_ct = max(ct_zscores, key=ct_zscores.get)
-            rows.append({
-                "module_idx": module_i, "gene": gene, "top_celltype": top_ct,
-                "z_score": ct_zscores[top_ct],
-                "global_mean_attn": float(g_mean[g_idx]),
-                "ct_mean_attn": float(ct_mean[top_ct][module_i][g_idx]),
-            })
-    return pd.DataFrame(rows)
-
-
-def compute_cohens_d(
-    global_mean: Dict[int, np.ndarray],
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    ct_std: Dict[str, Dict[int, np.ndarray]],
-    ct_count: Dict[str, int],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-    n_cells: int,
-    min_cells: int = 30,
-) -> pd.DataFrame:
-    """Cohen's d effect size for cell-type specificity."""
-    valid_cts = [ct for ct, n in ct_count.items() if n >= min_cells]
-    rows = []
-    for module_i in range(num_modules):
-        genes = module_gene_names[module_i]
-        for g_idx, gene in enumerate(genes):
-            best_d = -np.inf
-            best_ct = valid_cts[0]
-            for ct in valid_cts:
-                n_ct = ct_count[ct]
-                n_other = n_cells - n_ct
-                ct_m = float(ct_mean[ct][module_i][g_idx])
-                ct_s = float(ct_std[ct][module_i][g_idx])
-                other_m = (float(global_mean[module_i][g_idx]) * n_cells - ct_m * n_ct) / max(n_other, 1)
-                pooled = max(ct_s, 1e-10)
-                d = (ct_m - other_m) / pooled
-                if d > best_d:
-                    best_d = d
-                    best_ct = ct
-            rows.append({"module_idx": module_i, "gene": gene, "top_celltype": best_ct, "cohens_d": best_d})
-    return pd.DataFrame(rows)
-
-
-def compute_welch_ttest(
-    global_mean: Dict[int, np.ndarray],
-    global_std: Dict[int, np.ndarray],
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    ct_std: Dict[str, Dict[int, np.ndarray]],
-    ct_count: Dict[str, int],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-    n_cells: int,
-    min_cells: int = 30,
-) -> pd.DataFrame:
-    """Welch's t-test from streaming statistics with BH FDR correction."""
-    valid_cts = [ct for ct, n in ct_count.items() if n >= min_cells]
-    rows = []
-    for module_i in range(num_modules):
-        genes = module_gene_names[module_i]
-        for g_idx, gene in enumerate(genes):
-            best_t = -np.inf
-            best_ct = valid_cts[0]
-            best_pval = 1.0
-            for ct in valid_cts:
-                n_ct = ct_count[ct]
-                ct_m = float(ct_mean[ct][module_i][g_idx])
-                ct_v = float(ct_std[ct][module_i][g_idx]) ** 2
-                g_m = float(global_mean[module_i][g_idx])
-                g_v = float(global_std[module_i][g_idx]) ** 2
-                se = np.sqrt(ct_v / max(n_ct, 1) + g_v / max(n_cells, 1))
-                if se < 1e-12:
-                    continue
-                t_stat = (ct_m - g_m) / se
-                num = (ct_v / n_ct + g_v / n_cells) ** 2
-                denom = (ct_v / n_ct) ** 2 / max(n_ct - 1, 1) + (g_v / n_cells) ** 2 / max(n_cells - 1, 1)
-                df = num / max(denom, 1e-12)
-                pval = float(t_dist.sf(abs(t_stat), df) * 2)
-                if t_stat > best_t:
-                    best_t = t_stat
-                    best_ct = ct
-                    best_pval = pval
-            rows.append({"module_idx": module_i, "gene": gene, "top_celltype": best_ct,
-                         "t_stat": best_t, "pval": best_pval})
-
-    df = pd.DataFrame(rows)
-    # BH FDR correction
-    pvals = df["pval"].values
-    n = len(pvals)
-    sorted_idx = np.argsort(pvals)
-    sorted_pvals = pvals[sorted_idx]
-    fdr = np.minimum(1.0, sorted_pvals * n / (np.arange(n) + 1))
-    for i in range(n - 2, -1, -1):
-        fdr[i] = min(fdr[i], fdr[i + 1])
-    fdr_all = np.empty(n)
-    fdr_all[sorted_idx] = fdr
-    df["fdr"] = fdr_all
-    return df
-
-
-def compute_attention_expression_corr(
-    global_mean: Dict[int, np.ndarray],
-    expr_mean: Dict[int, np.ndarray],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-) -> Dict[str, float]:
-    """Correlation between mean attention and mean expression across genes."""
-    all_attn = []
-    all_expr = []
-    for module_i in range(num_modules):
-        all_attn.extend(global_mean[module_i].tolist())
-        all_expr.extend(expr_mean[module_i].tolist())
-    all_attn = np.array(all_attn)
-    all_expr = np.array(all_expr)
-    mask = (all_attn > 0) & (all_expr > 0)
-    if mask.sum() > 10:
-        pearson_r = float(np.corrcoef(all_attn[mask], all_expr[mask])[0, 1])
-        sp_r, sp_p = spearmanr(all_attn[mask], all_expr[mask])
-    else:
-        pearson_r, sp_r, sp_p = 0.0, 0.0, 1.0
-    return {"pearson_r": pearson_r, "spearman_r": float(sp_r),
-            "spearman_pval": float(sp_p), "n_genes": int(mask.sum())}
-
-
-def compute_random_baseline(
-    global_mean: Dict[int, np.ndarray],
-    num_modules: int,
-) -> Dict[str, float]:
-    """Compare learned attention to uniform baseline."""
-    kl_divs = []
-    learned_entropies = []
-    for module_i in range(num_modules):
-        attn = global_mean[module_i].astype(np.float64)
-        attn = attn / max(attn.sum(), 1e-10)
-        g_size = len(attn)
-        uniform = np.ones(g_size) / g_size
-        mask = attn > 1e-12
-        kl = float(np.sum(attn[mask] * np.log(attn[mask] / uniform[mask])))
-        kl_divs.append(kl)
-        h_learned = -float(np.sum(attn[mask] * np.log(attn[mask])))
-        h_max = np.log(g_size)
-        learned_entropies.append(h_learned / max(h_max, 1e-10))
-    return {
-        "mean_kl_divergence": float(np.mean(kl_divs)),
-        "std_kl_divergence": float(np.std(kl_divs)),
-        "mean_normalized_entropy": float(np.mean(learned_entropies)),
-        "entropy_reduction": float(1.0 - np.mean(learned_entropies)),
-        "n_modules": num_modules,
-    }
-
-
-def compute_attention_concentration(
-    global_mean: Dict[int, np.ndarray],
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    ct_count: Dict[str, int],
-    num_modules: int,
-    min_cells: int = 30,
-) -> pd.DataFrame:
-    """Gini coefficient per module measuring attention selectivity."""
-    def gini(arr):
-        arr = np.sort(arr.flatten())
-        n = len(arr)
-        if n == 0 or arr.sum() == 0:
-            return 0.0
-        index = np.arange(1, n + 1)
-        return float((2 * np.sum(index * arr) - (n + 1) * np.sum(arr)) / (n * np.sum(arr)))
-
-    valid_cts = [ct for ct, n in ct_count.items() if n >= min_cells]
-    rows = []
-    for module_i in range(num_modules):
-        g_attn = global_mean[module_i]
-        g_gini = gini(g_attn)
-        ct_ginis = [gini(ct_mean[ct][module_i]) for ct in valid_cts]
-        rows.append({
-            "module_idx": module_i, "global_gini": g_gini,
-            "mean_ct_gini": float(np.mean(ct_ginis)) if ct_ginis else 0.0,
-            "module_size": len(g_attn),
-        })
-    return pd.DataFrame(rows)
-
-
-def compute_marker_rank_within_celltype(
-    ct_mean: Dict[str, Dict[int, np.ndarray]],
-    ct_count: Dict[str, int],
-    module_gene_names: Dict[int, List[str]],
-    num_modules: int,
-    marker_dict: Dict[str, List[str]],
-    min_cells: int = 30,
-) -> pd.DataFrame:
-    """Rank marker genes within their expected cell type by attention."""
-    rows = []
-    for ct, markers in marker_dict.items():
-        if ct not in ct_count or ct_count[ct] < min_cells:
-            continue
-        gene_attn = {}
-        for module_i in range(num_modules):
-            genes = module_gene_names[module_i]
-            attn = ct_mean[ct][module_i]
-            for g_idx, gene in enumerate(genes):
-                if gene not in gene_attn or float(attn[g_idx]) > gene_attn[gene]:
-                    gene_attn[gene] = float(attn[g_idx])
-        sorted_genes = sorted(gene_attn.items(), key=lambda x: -x[1])
-        gene_rank = {g: rank + 1 for rank, (g, _) in enumerate(sorted_genes)}
-        n_total = len(gene_rank)
-        for marker in markers:
-            if marker in gene_rank:
-                rank = gene_rank[marker]
-                rows.append({
-                    "celltype": ct, "marker_gene": marker, "rank": rank,
-                    "percentile": rank / n_total * 100, "n_total_genes": n_total,
-                    "attention": gene_attn[marker],
-                })
-    return pd.DataFrame(rows)
-
-
-# ---------------------------------------------------------------------------
-# Marker gene validation
-# ---------------------------------------------------------------------------
-
-def validate_markers(
-    specificity_df: pd.DataFrame,
-    marker_dict: Dict[str, List[str]],
-    sort_col: str = "fold_change",
-) -> Dict[str, object]:
-    """Validate whether high-attention genes match known markers."""
-    results = {}
-    all_genes_in_model = set(specificity_df["gene"].unique())
-    for ct, markers in marker_dict.items():
-        markers_in_model = [g for g in markers if g in all_genes_in_model]
-        if not markers_in_model:
-            continue
-        ct_specific = specificity_df[specificity_df["top_celltype"] == ct].copy()
-        ct_specific = ct_specific.sort_values(sort_col, ascending=False).reset_index(drop=True)
-        marker_set = set(markers_in_model)
-        n_total = len(ct_specific)
-        if n_total == 0:
-            continue
-        top_50_genes = set(ct_specific.head(50)["gene"])
-        top_100_genes = set(ct_specific.head(100)["gene"])
-        hits_50 = len(marker_set & top_50_genes)
-        hits_100 = len(marker_set & top_100_genes)
-        marker_rows = ct_specific[ct_specific["gene"].isin(marker_set)]
-        if len(marker_rows) > 0:
-            marker_ranks = (marker_rows.index + 1).tolist()
-            mean_rank = np.mean(marker_ranks)
-            median_rank = np.median(marker_ranks)
-        else:
-            mean_rank = median_rank = n_total
-        results[ct] = {
-            "n_markers_in_model": len(markers_in_model),
-            "markers_in_model": markers_in_model,
-            "n_genes_assigned_to_ct": n_total,
-            "hits_in_top50": hits_50, "hits_in_top100": hits_100,
-            "precision_at_50": hits_50 / min(50, n_total),
-            "precision_at_100": hits_100 / min(100, n_total),
-            "marker_mean_rank": float(mean_rank),
-            "marker_median_rank": float(median_rank),
-            "recall_at_50": hits_50 / len(markers_in_model) if markers_in_model else 0,
-            "recall_at_100": hits_100 / len(markers_in_model) if markers_in_model else 0,
-        }
-    return results
-
-
-def validate_deg_correlation(
-    ct_dfs: Dict[str, pd.DataFrame],
-    deg_df: pd.DataFrame,
-) -> Dict[str, Dict[str, float]]:
-    """Compare attention ranking with DEG ranking per cell type."""
-    results = {}
-    for ct, attn_df in ct_dfs.items():
-        ct_deg = deg_df[deg_df["cell_type"] == ct].copy()
-        if len(ct_deg) == 0:
-            continue
-        merged = attn_df.merge(ct_deg[["gene", "log2fc"]], on="gene", how="inner")
-        if len(merged) < 10:
-            continue
-        rho, pval = spearmanr(merged["mean_attention"], merged["log2fc"].abs())
-        results[ct] = {"spearman_r": float(rho), "pval": float(pval), "n_genes": len(merged)}
-    return results
-
-
-# ---------------------------------------------------------------------------
-# Publication-quality visualization
-# ---------------------------------------------------------------------------
-
-def plot_global_importance_pub(global_df, specificity_zscore_df, top_k, output_dir):
-    """Fig1: Top-K genes with std error bars, colored by top cell type."""
-    setup_publication_style()
-    gene_stats = global_df.groupby("gene").agg(
-        mean_attn=("mean_attention", "mean"), std_attn=("std_attention", "mean"),
-    ).nlargest(top_k, "mean_attn")
-    gene_ct = specificity_zscore_df.drop_duplicates("gene").set_index("gene")["top_celltype"]
-    fig, ax = plt.subplots(figsize=(4.5, 5))
-    genes = gene_stats.index.tolist()
-    y_pos = np.arange(len(genes))
-    colors = [PANCREAS_COLORS.get(gene_ct.get(g, ""), "#999999") for g in genes]
-    ax.barh(y_pos, gene_stats["mean_attn"], xerr=gene_stats["std_attn"],
-            color=colors, edgecolor="none", height=0.7, capsize=2, error_kw={"linewidth": 0.8})
-    ax.set_yticks(y_pos)
-    ax.set_yticklabels(genes)
-    ax.set_xlabel("Mean attention weight")
-    ax.set_title("Top genes by global attention")
-    ax.invert_yaxis()
-    pub_savefig(fig, output_dir / "fig1_global_importance")
-    logger.info("Saved fig1_global_importance (PDF+PNG)")
-
-
-def plot_celltype_heatmap_pub(ct_dfs, ct_count, top_k, output_dir, min_cells=30):
-    """Fig2: Clustered heatmap of gene attention (z-score normalized)."""
-    setup_publication_style()
-    valid_cts = {ct for ct, n in ct_count.items() if n >= min_cells}
-    all_genes_attn = {}
-    for ct, df in ct_dfs.items():
-        if ct not in valid_cts:
-            continue
-        all_genes_attn[ct] = df.groupby("gene")["mean_attention"].mean()
-    combined = pd.DataFrame(all_genes_attn)
-    gene_var = combined.var(axis=1).nlargest(top_k).index
-    heatmap_data = combined.loc[gene_var]
-    row_mean = heatmap_data.mean(axis=1)
-    row_std = heatmap_data.std(axis=1).replace(0, 1)
-    heatmap_z = heatmap_data.sub(row_mean, axis=0).div(row_std, axis=0)
-    g = sns.clustermap(heatmap_z, cmap="RdBu_r", center=0, figsize=(6, 7),
-                       xticklabels=True, yticklabels=True,
-                       dendrogram_ratio=(0.1, 0.15), cbar_kws={"label": "Z-score"},
-                       linewidths=0.3)
-    g.ax_heatmap.set_xlabel("")
-    g.ax_heatmap.set_ylabel("")
-    g.fig.suptitle(f"Top {top_k} variable genes (z-scored attention)", y=1.02, fontsize=11)
-    pub_savefig(g.fig, output_dir / "fig2_celltype_heatmap")
-    logger.info("Saved fig2_celltype_heatmap (PDF+PNG)")
-
-
-def plot_specificity_volcano(welch_df, zscore_df, top_k, output_dir):
-    """Fig3: Volcano plot (z-score vs -log10 FDR)."""
-    setup_publication_style()
-    merged = zscore_df[["gene", "top_celltype", "z_score"]].drop_duplicates("gene")
-    welch_best = welch_df.loc[welch_df.groupby("gene")["t_stat"].idxmax()][["gene", "fdr"]]
-    merged = merged.merge(welch_best, on="gene", how="left")
-    merged["neg_log10_fdr"] = -np.log10(merged["fdr"].clip(lower=1e-300))
-    fig, ax = plt.subplots(figsize=(5.5, 5))
-    for ct in merged["top_celltype"].unique():
-        subset = merged[merged["top_celltype"] == ct]
-        color = PANCREAS_COLORS.get(ct, "#999999")
-        ax.scatter(subset["z_score"], subset["neg_log10_fdr"],
-                   c=color, s=8, alpha=0.5, label=ct, rasterized=True)
-    top_genes = merged.nlargest(top_k, "z_score")
-    for _, row in top_genes.iterrows():
-        ax.annotate(row["gene"], (row["z_score"], row["neg_log10_fdr"]),
-                    fontsize=6, ha="left", va="bottom")
-    ax.axhline(-np.log10(0.05), color="gray", linestyle="--", linewidth=0.6)
-    ax.axvline(2, color="gray", linestyle="--", linewidth=0.6)
-    ax.set_xlabel("Z-score (cell-type specificity)")
-    ax.set_ylabel(r"$-\log_{10}$(FDR)")
-    ax.set_title("Gene specificity volcano")
-    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=7, frameon=False, markerscale=2)
-    pub_savefig(fig, output_dir / "fig3_specificity_volcano")
-    logger.info("Saved fig3_specificity_volcano (PDF+PNG)")
-
-
-def plot_marker_dotplot(marker_rank_df, output_dir):
-    """Fig4: Dot plot of marker gene rank percentile per cell type."""
-    setup_publication_style()
-    if marker_rank_df.empty:
-        return
-    fig, ax = plt.subplots(figsize=(6, 4.5))
-    celltypes = marker_rank_df["celltype"].unique()
-    ct_positions = {ct: i for i, ct in enumerate(celltypes)}
-    for _, row in marker_rank_df.iterrows():
-        y = ct_positions[row["celltype"]]
-        size = max(10, 200 * (1 - row["percentile"] / 100))
-        color = PANCREAS_COLORS.get(row["celltype"], "#999999")
-        ax.scatter(row["percentile"], y, s=size, c=color, alpha=0.7,
-                   edgecolors="white", linewidth=0.3)
-        if row["percentile"] < 5:
-            ax.annotate(row["marker_gene"], (row["percentile"], y),
-                        fontsize=6, ha="left", va="center", xytext=(3, 0),
-                        textcoords="offset points")
-    ax.set_yticks(range(len(celltypes)))
-    ax.set_yticklabels(celltypes)
-    ax.set_xlabel("Rank percentile (lower = higher attention)")
-    ax.set_title("Marker gene rank within expected cell type")
-    ax.axvline(5, color="gray", linestyle="--", linewidth=0.6, alpha=0.5, label="Top 5%")
-    ax.legend(fontsize=7, frameon=False)
-    ax.set_xlim(-2, 105)
-    pub_savefig(fig, output_dir / "fig4_marker_dotplot")
-    logger.info("Saved fig4_marker_dotplot (PDF+PNG)")
-
-
-def plot_marker_recall_pub(marker_results, ct_count, output_dir, min_cells=30):
-    """Fig5: Marker recall bar chart with random expectation line."""
-    setup_publication_style()
-    if not marker_results:
-        return
-    cts = [ct for ct in marker_results.keys() if ct_count.get(ct, 0) >= min_cells]
-    if not cts:
-        return
-    recall_50 = [marker_results[ct]["recall_at_50"] for ct in cts]
-    recall_100 = [marker_results[ct]["recall_at_100"] for ct in cts]
-    x = np.arange(len(cts))
-    width = 0.35
-    fig, ax = plt.subplots(figsize=(max(5, len(cts) * 0.7), 4))
-    colors_50 = [PANCREAS_COLORS.get(ct, "#999999") for ct in cts]
-    ax.bar(x - width / 2, recall_50, width, color=colors_50, alpha=0.9, label="Recall@50")
-    ax.bar(x + width / 2, recall_100, width, color=colors_50, alpha=0.5, label="Recall@100")
-    # Random expectation
-    for i, ct in enumerate(cts):
-        n_genes = marker_results[ct]["n_genes_assigned_to_ct"]
-        if n_genes > 0:
-            random_50 = min(50, n_genes) / n_genes
-            ax.hlines(random_50, i - 0.4, i + 0.4, colors="gray", linestyles=":", linewidth=0.8)
-    ax.set_xticks(x)
-    ax.set_xticklabels(cts, rotation=45, ha="right")
-    ax.set_ylabel("Recall")
-    ax.set_title("Marker gene recall by cell type")
-    ax.legend(fontsize=7, frameon=False)
-    pub_savefig(fig, output_dir / "fig5_marker_recall")
-    logger.info("Saved fig5_marker_recall (PDF+PNG)")
-
-
-def plot_attention_vs_expression(global_mean, expr_mean, module_gene_names, num_modules, output_dir):
-    """Fig6: Scatter of mean attention vs mean expression."""
-    setup_publication_style()
-    all_attn, all_expr = [], []
-    for module_i in range(num_modules):
-        all_attn.extend(global_mean[module_i].tolist())
-        all_expr.extend(expr_mean[module_i].tolist())
-    all_attn = np.array(all_attn)
-    all_expr = np.array(all_expr)
-    mask = (all_attn > 0) & (all_expr > 0)
-    fig, ax = plt.subplots(figsize=(4.5, 4.5))
-    ax.scatter(all_expr[mask], all_attn[mask], s=3, alpha=0.3, c="#377EB8", rasterized=True)
-    r = np.corrcoef(all_expr[mask], all_attn[mask])[0, 1]
-    ax.set_xlabel("Mean expression")
-    ax.set_ylabel("Mean attention weight")
-    ax.set_title("Attention vs expression")
-    ax.text(0.05, 0.95, f"Pearson r = {r:.3f}", transform=ax.transAxes,
-            fontsize=9, va="top", ha="left")
-    pub_savefig(fig, output_dir / "fig6_attn_vs_expression")
-    logger.info("Saved fig6_attn_vs_expression (PDF+PNG)")
-
-
-def plot_concentration_pub(concentration_df, baseline_stats, output_dir):
-    """Fig7: Gini coefficient distribution (learned vs uniform)."""
-    setup_publication_style()
-    fig, axes = plt.subplots(1, 2, figsize=(8, 3.5))
-    # Panel a: Gini histogram
-    ax = axes[0]
-    ax.hist(concentration_df["global_gini"], bins=30, alpha=0.7, color="#377EB8", label="Global")
-    ax.hist(concentration_df["mean_ct_gini"], bins=30, alpha=0.7, color="#E41A1C", label="Per cell-type")
-    ax.axvline(0, color="gray", linestyle="--", linewidth=0.6)
-    ax.set_xlabel("Gini coefficient")
-    ax.set_ylabel("Number of modules")
-    ax.set_title("Attention concentration")
-    ax.legend(fontsize=7, frameon=False)
-    # Panel b: Entropy comparison
-    ax = axes[1]
-    labels = ["Learned", "Uniform"]
-    values = [baseline_stats["mean_normalized_entropy"], 1.0]
-    bars = ax.bar(labels, values, color=["#377EB8", "#CCCCCC"], width=0.5)
-    ax.set_ylabel("Normalized entropy")
-    ax.set_title(f"Entropy reduction: {baseline_stats['entropy_reduction']:.1%}")
-    ax.set_ylim(0, 1.1)
-    plt.tight_layout()
-    pub_savefig(fig, output_dir / "fig7_concentration")
-    logger.info("Saved fig7_concentration (PDF+PNG)")
-
-
-def plot_baseline_kl(baseline_stats, output_dir):
-    """Fig8: KL divergence summary."""
-    setup_publication_style()
-    fig, ax = plt.subplots(figsize=(3.5, 3.5))
-    ax.bar(["KL(learned||uniform)"], [baseline_stats["mean_kl_divergence"]],
-           yerr=[baseline_stats["std_kl_divergence"]], color="#377EB8",
-           capsize=5, width=0.4)
-    ax.set_ylabel("KL divergence (nats)")
-    ax.set_title("Deviation from uniform attention")
-    pub_savefig(fig, output_dir / "fig8_baseline_kl")
-    logger.info("Saved fig8_baseline_kl (PDF+PNG)")
-
-
-def plot_deg_correlation_pub(deg_results, output_dir):
-    """Fig9: DEG correlation with significance."""
-    setup_publication_style()
-    if not deg_results:
-        return
-    cts = sorted(deg_results.keys(), key=lambda c: deg_results[c]["spearman_r"], reverse=True)
-    rhos = [deg_results[ct]["spearman_r"] for ct in cts]
-    pvals = [deg_results[ct]["pval"] for ct in cts]
-    fig, ax = plt.subplots(figsize=(max(5, len(cts) * 0.6), 4))
-    colors = ["#377EB8" if r > 0 else "#E41A1C" for r in rhos]
-    bars = ax.bar(range(len(cts)), rhos, color=colors, width=0.6)
-    for i, (bar, p) in enumerate(zip(bars, pvals)):
-        if p < 0.001:
-            ax.text(i, bar.get_height() + 0.02, "***", ha="center", fontsize=7)
-        elif p < 0.01:
-            ax.text(i, bar.get_height() + 0.02, "**", ha="center", fontsize=7)
-        elif p < 0.05:
-            ax.text(i, bar.get_height() + 0.02, "*", ha="center", fontsize=7)
-    ax.set_xticks(range(len(cts)))
-    ax.set_xticklabels(cts, rotation=45, ha="right")
-    ax.set_ylabel("Spearman r")
-    ax.set_title("Attention vs DEG correlation")
-    ax.axhline(0, color="gray", linewidth=0.5)
-    pub_savefig(fig, output_dir / "fig9_deg_correlation")
-    logger.info("Saved fig9_deg_correlation (PDF+PNG)")
-
-
-# ---------------------------------------------------------------------------
-# Enrichment analysis (optional)
-# ---------------------------------------------------------------------------
-
-def run_enrichment_analysis(
-    global_df: pd.DataFrame, top_k: int, output_dir: Path, organism: str = "human",
-) -> Optional[pd.DataFrame]:
-    """Run GO enrichment on top-K attention genes using gseapy."""
-    try:
-        import gseapy as gp
-    except ImportError:
-        logger.warning("gseapy not installed, skipping enrichment. Install: pip install gseapy")
-        return None
-    top_genes = global_df.groupby("gene")["mean_attention"].mean().nlargest(top_k * 5).index.tolist()
-    logger.info("Running GO enrichment on top %d genes...", len(top_genes))
-    try:
-        enr = gp.enrichr(
-            gene_list=top_genes,
-            gene_sets=["GO_Biological_Process_2021", "KEGG_2021_Human"],
-            organism=organism,
-            outdir=str(output_dir / "enrichment"),
-            no_plot=True,
-        )
-        results_df = enr.results
-        results_df.to_csv(output_dir / "enrichment_results.csv", index=False)
-        logger.info("Enrichment results saved (%d terms)", len(results_df))
-        return results_df
-    except Exception as e:
-        logger.warning("Enrichment analysis failed: %s", e)
-        return None
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-def main():
-    args = parse_args()
-    seed_all(args.seed)
-
-    _script_dir = Path(__file__).resolve().parent
-    for attr in ("performer_checkpoint", "config_pkl", "module_data_dir", "gene_list_path",
-                 "gene_mapping_path", "gene2vec_data_path", "gene2ens_file", "data_path"):
-        val = getattr(args, attr, None)
-        if val is not None and not Path(val).is_absolute():
-            setattr(args, attr, str((_script_dir / val).resolve()))
-
-    output_dir = ensure_dir(Path(args.output_dir))
-    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-    logger.info("Device: %s", device)
-
-    # --- Load model ---
-    logger.info("Loading backbone from %s", args.performer_checkpoint)
-    backbone, biomod_data, model_config = load_backbone(args, device)
-    backbone.eval()
-
-    # --- Build module gene name mapping ---
-    model_genes_file = args.gene_list_path
-    with open(model_genes_file) as f:
-        model_genes = [line.strip() for line in f if line.strip()]
-
-    gene2ens_df = pd.read_csv(args.gene2ens_file)
-    ens_to_symbol = dict(zip(gene2ens_df["Ensembl gene ID"], gene2ens_df.iloc[:, 0]))
-
-    module_sets = backbone.module_embed.module_sets
-    num_modules = len(module_sets)
-    module_sizes = [len(ps) for ps in module_sets]
-
-    module_gene_names: Dict[int, List[str]] = {}
-    for module_i, gene_idx in enumerate(module_sets):
-        names = []
-        for idx in gene_idx.cpu().numpy():
-            ens_id = model_genes[idx]
-            symbol = ens_to_symbol.get(ens_id, ens_id)
-            names.append(symbol)
-        module_gene_names[module_i] = names
-
-    logger.info("Modules: %d, total genes: %d", num_modules, sum(module_sizes))
-
-    # --- Load and preprocess data ---
-    logger.info("Loading data from %s", args.data_path)
-    preprocessor = DataPreprocessor(
-        gene2ens_file=args.gene2ens_file,
-        gene2ens_column=args.gene2ens_column,
-        model_genes_file=model_genes_file,
-    )
-    preprocessor.load_gene_mapping()
-    preprocessor.load_model_genes()
-
-    gene_name_col = args.gene_name_column if args.gene_name_column != "auto" else None
-    adata, reindex_state = preprocessor.preprocess_adata(
-        adata_path=args.data_path,
-        max_cells=args.max_cells,
-        dataset_name=args.dataset_name,
-        gene_name_column=gene_name_col,
-    )
-    cell_labels = adata.obs[args.celltype_column].values.astype(str)
-    logger.info("Data: %d cells, %d cell types", len(adata), len(np.unique(cell_labels)))
-
-    # --- DataLoader ---
-    if hasattr(adata.X, "toarray"):
-        X = adata.X.toarray().astype(np.float32)
-    else:
-        X = np.array(adata.X, dtype=np.float32)
-
-    dataset = MatrixRowDataset(X)
-    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
-                        num_workers=args.num_workers, pin_memory=True)
-
-    # --- Extract attention (streaming) ---
-    logger.info("Extracting module attention (streaming mode)...")
-    (global_mean, global_std, ct_mean, ct_std, ct_count, n_cells,
-     expr_mean, attn_expr_prod_sum) = extract_attention_streaming(
-        backbone, loader, preprocessor, reindex_state, device,
-        cell_labels, num_modules, module_sizes,
-    )
-
-    # --- Build DataFrames ---
-    logger.info("Building analysis DataFrames...")
-    global_df = build_global_importance_df(global_mean, global_std, module_gene_names, num_modules)
-    ct_dfs = build_celltype_dfs(ct_mean, ct_count, module_gene_names, num_modules)
-    specificity_df = compute_specificity_from_means(global_mean, ct_mean, module_gene_names, num_modules)
-
-    # --- New metrics ---
-    logger.info("Computing z-score specificity...")
-    zscore_df = compute_specificity_zscore(
-        global_mean, global_std, ct_mean, ct_count, module_gene_names, num_modules, args.min_cells)
-
-    logger.info("Computing Cohen's d...")
-    cohens_d_df = compute_cohens_d(
-        global_mean, ct_mean, ct_std, ct_count, module_gene_names, num_modules, n_cells, args.min_cells)
-
-    logger.info("Computing Welch's t-test...")
-    welch_df = compute_welch_ttest(
-        global_mean, global_std, ct_mean, ct_std, ct_count, module_gene_names,
-        num_modules, n_cells, args.min_cells)
-
-    logger.info("Computing attention-expression correlation...")
-    attn_expr_corr = compute_attention_expression_corr(
-        global_mean, expr_mean, module_gene_names, num_modules)
-
-    logger.info("Computing random baseline comparison...")
-    baseline_stats = compute_random_baseline(global_mean, num_modules)
-
-    logger.info("Computing attention concentration...")
-    concentration_df = compute_attention_concentration(
-        global_mean, ct_mean, ct_count, num_modules, args.min_cells)
-
-    # --- Marker validation ---
-    marker_dict = PANCREAS_MARKERS
-    if args.marker_db:
-        with open(args.marker_db) as f:
-            marker_dict = json.load(f)
-
-    logger.info("Computing marker ranks within cell types...")
-    marker_rank_df = compute_marker_rank_within_celltype(
-        ct_mean, ct_count, module_gene_names, num_modules, marker_dict, args.min_cells)
-
-    marker_results_zscore = validate_markers(zscore_df, marker_dict, sort_col="z_score")
-    marker_results_fc = validate_markers(specificity_df, marker_dict, sort_col="fold_change")
-
-    # --- Save CSVs ---
-    logger.info("Saving CSV results...")
-    global_df.to_csv(output_dir / "global_gene_importance.csv", index=False)
-    specificity_df.to_csv(output_dir / "gene_specificity.csv", index=False)
-    zscore_df.to_csv(output_dir / "gene_specificity_zscore.csv", index=False)
-    cohens_d_df.to_csv(output_dir / "gene_specificity_cohens_d.csv", index=False)
-    welch_df.to_csv(output_dir / "gene_specificity_welch.csv", index=False)
-    concentration_df.to_csv(output_dir / "attention_concentration.csv", index=False)
-    if not marker_rank_df.empty:
-        marker_rank_df.to_csv(output_dir / "marker_ranks_within_celltype.csv", index=False)
-    for ct, df in ct_dfs.items():
-        safe_name = ct.replace("/", "_").replace(" ", "_")
-        df.to_csv(output_dir / f"celltype_{safe_name}_attention.csv", index=False)
-
-    with open(output_dir / "marker_validation_zscore.json", "w") as f:
-        json.dump(marker_results_zscore, f, indent=2, ensure_ascii=False)
-    with open(output_dir / "marker_validation.json", "w") as f:
-        json.dump(marker_results_fc, f, indent=2, ensure_ascii=False)
-
-    # --- DEG correlation ---
-    deg_results = None
-    if args.deg_csv:
-        deg_df = pd.read_csv(args.deg_csv)
-        deg_results = validate_deg_correlation(ct_dfs, deg_df)
-        with open(output_dir / "deg_correlation.json", "w") as f:
-            json.dump(deg_results, f, indent=2)
-
-    # --- Publication figures ---
-    logger.info("Generating publication figures...")
-    plot_global_importance_pub(global_df, zscore_df, args.top_k_genes, output_dir)
-    plot_celltype_heatmap_pub(ct_dfs, ct_count, args.top_k_genes, output_dir, args.min_cells)
-    plot_specificity_volcano(welch_df, zscore_df, args.top_k_genes, output_dir)
-    plot_marker_dotplot(marker_rank_df, output_dir)
-    plot_marker_recall_pub(marker_results_zscore, ct_count, output_dir, args.min_cells)
-    plot_attention_vs_expression(global_mean, expr_mean, module_gene_names, num_modules, output_dir)
-    plot_concentration_pub(concentration_df, baseline_stats, output_dir)
-    plot_baseline_kl(baseline_stats, output_dir)
-    if deg_results:
-        plot_deg_correlation_pub(deg_results, output_dir)
-
-    # --- Enrichment ---
-    run_enrichment_analysis(global_df, args.top_k_genes, output_dir)
-
-    # --- Summary ---
-    filtered_cts = [ct for ct, n in ct_count.items() if n < args.min_cells]
-    summary = {
-        "n_cells": n_cells,
-        "n_modules": num_modules,
-        "n_celltypes": len(ct_count),
-        "celltypes": {ct: int(n) for ct, n in ct_count.items()},
-        "min_cells_filter": args.min_cells,
-        "filtered_celltypes": filtered_cts,
-        "top_10_global_genes": global_df.groupby("gene")["mean_attention"].mean()
-            .nlargest(10).to_dict(),
-        "attn_expr_correlation": attn_expr_corr,
-        "baseline_comparison": baseline_stats,
-        "mean_gini_global": float(concentration_df["global_gini"].mean()),
-        "mean_gini_celltype": float(concentration_df["mean_ct_gini"].mean()),
-    }
-    with open(output_dir / "summary.json", "w") as f:
-        json.dump(summary, f, indent=2, ensure_ascii=False)
-
-    logger.info("Done. Results saved to %s", output_dir)
-
-
-if __name__ == "__main__":
-    main()
-

BioMod/interpretability/biomod_geometry.py

@@ -1,655 +0,0 @@
-"""Experiment 1 — Intrinsic geometry of module embeddings.
-
-Question
---------
-Does scLong's encoder organize its module latent space by pathway family
-without any explicit family-level supervision,  If yes, this is a structural
-property gene-level baselines (GenePT / scGPT) cannot exhibit because they
-do not have a module tier.
-
-Pipeline
---------
-1. Load pretrained backbone (forward with mask_ratio=0).
-2. Iterate over a population of cells, accumulate the per-module encoder output
-   in a streaming average -> [num_modules, base_dim].
-3. Parse each module's family from merged_biomods.json keys
-   (Hallmarks / Reactome / Hotspot / Processed) and, where applicable, the
-   MSigDB hallmark process subcategory (immune / metabolic / signaling / ...).
-4. Quantify family separation:
-       * silhouette score (cosine + euclidean)
-       * kNN purity at k=10
-       * inter-family / intra-family mean-distance ratio
-   Compare each metric against a permutation null (shuffle family labels K
-   times) and report empirical p-values.
-5. Visualise with UMAP and PCA, coloured by family / hallmark process.
-6. Save embeddings, metadata, metrics.json, and figures.
-
-Data modes
-----------
---data_mode h5ad           Generic h5ad. Goes through DataPreprocessor and
-                           reindex_tensor_to_model_space. Works for
-                           cell_annotation / zero_shot / any annotated h5ad.
---data_mode pretrain_npz   Per-cell npz directory (cell0.npz, cell1.npz, ...),
-                           already in the 27874-dim model gene space. Requires
-                           --cell_num. Skips reindexing.
---data_mode pretrain_h5ad  h5ad whose var.index is already the 27874-dim
-                           model gene list (selected_genes_27k.txt). Skips
-                           reindexing for speed.
-
-DDP
----
-Supports multi-GPU via torchrun. When WORLD_SIZE>1, each rank processes a
-shard of cells; embeddings are all_reduced before metrics/plots. Single-GPU
-and CPU work unchanged.
-
-CLI
----
-python BioMod/interpretability/biomod_geometry.py \\
-    --biomod_checkpoint <ckpt.pth> \\
-    --config_pkl <config.pkl> \\
-    --data_path <path.h5ad> \\
-    --data_mode h5ad \\
-    --output_dir ./interp/geometry/<run> \\
-    --max_cells 5000
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-import torch
-import torch.distributed as dist
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-from sklearn.decomposition import PCA
-from sklearn.metrics import silhouette_score
-from sklearn.metrics.pairwise import cosine_distances, euclidean_distances
-from sklearn.neighbors import NearestNeighbors
-
-# Make sibling modules importable when launched as a script
-HERE = Path(__file__).resolve()
-BIOMOD_SOURCE = HERE.parents[1]
-REPO_ROOT = HERE.parents[2]
-ASSET_ROOT = REPO_ROOT / "model_dependencies"
-for p in (str(BIOMOD_SOURCE), str(REPO_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from cell_annotation.biomod_finetune import (  # noqa: E402
-    DataPreprocessor,
-    get_rank,
-    get_world_size,
-    init_distributed_mode,
-    is_dist_initialized,
-    is_main_process,
-    load_backbone,
-    reduce_sum,
-    resolve_path as ca_resolve_path,
-    seed_all,
-)
-
-from interpretability.utils import (  # noqa: E402
-    MatrixRowDataset,
-    NpzCellDataset,
-    build_biomod_metadata,
-    ensure_dir,
-    load_biomod_names,
-)
-
-logger = logging.getLogger(__name__)
-
-
-# ---------- argparse ----------
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-
-    # Pretrained backbone (mirrors cell_annotation.biomod_finetune)
-    parser.add_argument("--biomod_checkpoint", "--performer_checkpoint", type=str, required=True, dest='performer_checkpoint')
-    parser.add_argument("--config_pkl", type=str, required=True)
-    parser.add_argument("--biomod_data_dir", "--module_data_dir", type=str, default=str(ASSET_ROOT / "biomod_data_all"), dest='module_data_dir')
-    parser.add_argument("--biomod_json", "--module_json", type=str, default=None,
-                        help="Path to merged_biomods.json. Defaults to <module_data_dir>/merged_biomods.json.", dest='module_json')
-    parser.add_argument("--gene_list_path", type=str, default=str(ASSET_ROOT / "selected_genes_27k.txt"))
-    parser.add_argument("--gene_mapping_path", type=str, default=str(ASSET_ROOT / "gene2ens.csv"))
-    parser.add_argument("--gene2vec_data_path", type=str, default=str(ASSET_ROOT / "selected_gene2vec_27k.npy"))
-    parser.add_argument("--gene2ens_file", type=str, default=str(ASSET_ROOT / "gene2ens.csv"))
-    parser.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
-    parser.add_argument("--shuffle_biomod_genes", "--shuffle_module_genes", action="store_true",
-                        help="Randomly reshuffle gene-to-module assignment (control / ablation).", dest='shuffle_module_genes')
-    parser.add_argument("--shuffle_seed", type=int, default=2021)
-
-    # Data
-    parser.add_argument("--data_path", type=str, required=True,
-                        help="h5ad file path, or npz directory for pretrain_npz mode.")
-    parser.add_argument("--data_mode", type=str, default="h5ad",
-                        choices=["h5ad", "pretrain_h5ad", "pretrain_npz"],
-                        help="h5ad: generic, goes through DataPreprocessor reindex. "
-                             "pretrain_h5ad: var.index is already the 27874 model gene list, skip reindex. "
-                             "pretrain_npz: per-cell npz directory (cell0.npz, ...), skip reindex.")
-    parser.add_argument("--cell_num", type=int, default=None,
-                        help="Total cell count for pretrain_npz sequential naming (cell0.npz..cellN.npz).")
-    parser.add_argument("--gene_name_column", type=str, default="auto",
-                        help="Which adata.var column holds gene symbols. 'auto' / 'index' / column name. "
-                             "Only used when data_mode=h5ad.")
-    parser.add_argument("--dataset_name", type=str, default="interpretability")
-    parser.add_argument("--min_genes", type=int, default=0)
-    parser.add_argument("--min_cells", type=int, default=0)
-    parser.add_argument("--max_cells", type=int, default=None,
-                        help="Optional subsample for inference speed. None = use all cells.")
-
-    # Inference
-    parser.add_argument("--device", type=str, default="cuda")
-    parser.add_argument("--batch_size", type=int, default=8)
-    parser.add_argument("--num_workers", type=int, default=0)
-    parser.add_argument("--seed", type=int, default=2021)
-
-    # Geometry / metrics
-    parser.add_argument("--knn_k", type=int, default=10,
-                        help="k for kNN-purity metric.")
-    parser.add_argument("--null_permutations", type=int, default=200,
-                        help="Number of label-permutation samples for null distribution.")
-    parser.add_argument("--include_other_in_metrics", action="store_true",
-                        help="If false (default), exclude modules whose family is 'Other' from metric computation.")
-
-    # UMAP
-    parser.add_argument("--umap_n_neighbors", type=int, default=15)
-    parser.add_argument("--umap_min_dist", type=float, default=0.3)
-
-    # Output
-    parser.add_argument("--output_dir", type=str, required=True)
-    parser.add_argument("--run_name", type=str, default="")
-
-    return parser.parse_args()
-
-
-# ---------- dataset construction ----------
-
-def build_inference_loader(
-    args: argparse.Namespace,
-    preprocessor: Optional[DataPreprocessor],
-    use_ddp: bool,
-    rank: int,
-    world_size: int,
-) -> Tuple[DataLoader, Optional[dict], int]:
-    """Build a DataLoader yielding raw [G_input] expression rows and a reindex_state.
-
-    Returns: (loader, reindex_state_or_None, n_cells_total)
-    """
-    if args.data_mode == "h5ad":
-        import anndata as ad
-        assert preprocessor is not None
-        adata, reindex_state = preprocessor.preprocess_adata(
-            adata_path=args.data_path,
-            min_genes=args.min_genes,
-            min_cells=args.min_cells,
-            max_cells=args.max_cells,
-            dataset_name=args.dataset_name,
-            gene_name_column=args.gene_name_column,
-        )
-        dataset = MatrixRowDataset(adata.X)
-        n_cells = adata.n_obs
-    elif args.data_mode == "pretrain_h5ad":
-        import anndata as ad
-        logger.info("Loading pretrain-mode h5ad %s", args.data_path)
-        adata = ad.read_h5ad(args.data_path)
-        if args.max_cells is not None and adata.n_obs > args.max_cells:
-            rng = np.random.default_rng(args.seed)
-            chosen = np.sort(rng.choice(adata.n_obs, size=args.max_cells, replace=False))
-            adata = adata[chosen].copy()
-        dataset = MatrixRowDataset(adata.X)
-        n_cells = adata.n_obs
-        reindex_state = None
-    elif args.data_mode == "pretrain_npz":
-        dataset = NpzCellDataset(
-            data_dir=args.data_path,
-            cell_num=args.cell_num,
-            max_cells=args.max_cells,
-            seed=args.seed,
-        )
-        n_cells = len(dataset)
-        reindex_state = None
-    else:
-        raise ValueError(f"Unsupported data_mode: {args.data_mode}")
-
-    if use_ddp:
-        sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank,
-                                     shuffle=False, drop_last=False)
-    else:
-        sampler = None
-
-    loader = DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-        sampler=sampler,
-        num_workers=args.num_workers,
-        pin_memory=False,
-    )
-    logger.info("Built inference loader: %d cells total, batch_size=%d, ddp=%s",
-                n_cells, args.batch_size, use_ddp)
-    return loader, reindex_state, n_cells
-
-
-# ---------- inference: cell-averaged module embeddings ----------
-
-@torch.no_grad()
-def extract_cell_averaged_biomod_embeddings(
-    backbone,
-    loader: DataLoader,
-    preprocessor: Optional[DataPreprocessor],
-    reindex_state: Optional[dict],
-    device: torch.device,
-    num_modules: int,
-    base_dim: int,
-    expected_input_genes: int,
-    use_ddp: bool,
-) -> Tuple[np.ndarray, int]:
-    """Stream forward over the loader and accumulate the cell-mean module embedding.
-
-    Under DDP each rank processes its shard; sum_emb and n_cells are all_reduced
-    before computing the mean.
-
-    Returns: (mean_embedding [num_modules, base_dim] float32, n_cells_total)
-    """
-    backbone.eval()
-    sum_emb = torch.zeros(num_modules, base_dim, device=device, dtype=torch.float64)
-    n_cells = 0
-
-    for batch_idx, x in enumerate(loader):
-        x = x.to(device, non_blocking=False)
-        if reindex_state is not None:
-            x_model = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
-        else:
-            if x.shape[1] != expected_input_genes:
-                raise ValueError(
-                    f"pretrain mode expects {expected_input_genes} genes per cell, got {x.shape[1]}. "
-                    f"Use --data_mode h5ad (with reindex) for arbitrary gene spaces."
-                )
-            x_model = x
-
-        out = backbone(x_model, mask_ratio=0.0, return_encodings=True)
-        enc = out["module_encodings"]  # (B, num_modules, D)
-        if enc.shape[1] != num_modules or enc.shape[2] != base_dim:
-            raise RuntimeError(
-                f"Encoder output shape mismatch: got {tuple(enc.shape)}, "
-                f"expected (*, {num_modules}, {base_dim}). "
-                "Module order between merged_biomods.json and the loaded backbone may have drifted."
-            )
-        sum_emb += enc.to(torch.float64).sum(dim=0)
-        n_cells += enc.shape[0]
-        if (batch_idx + 1) % 50 == 0:
-            logger.info("  forward batch %d (%d cells on this rank so far)", batch_idx + 1, n_cells)
-
-    if use_ddp:
-        dist.all_reduce(sum_emb, op=dist.ReduceOp.SUM)
-        n_cells_tensor = torch.tensor(n_cells, device=device, dtype=torch.long)
-        dist.all_reduce(n_cells_tensor, op=dist.ReduceOp.SUM)
-        n_cells = int(n_cells_tensor.item())
-
-    mean_emb = (sum_emb / max(n_cells, 1)).cpu().numpy().astype(np.float32)
-    return mean_emb, n_cells
-
-
-# ---------- structure metrics ----------
-
-def _label_to_int(labels: List[str]) -> np.ndarray:
-    uniq = sorted(set(labels))
-    mapping = {lab: i for i, lab in enumerate(uniq)}
-    return np.asarray([mapping[lab] for lab in labels], dtype=np.int64)
-
-
-def _knn_purity(emb: np.ndarray, labels: np.ndarray, k: int, metric: str = "cosine") -> float:
-    n = emb.shape[0]
-    k_use = min(k, n - 1)
-    nn = NearestNeighbors(n_neighbors=k_use + 1, metric=metric)
-    nn.fit(emb)
-    _, idx = nn.kneighbors(emb)
-    # idx[:, 0] == self; drop it
-    neighbor_labels = labels[idx[:, 1:]]
-    own_labels = labels[:, None]
-    purity = (neighbor_labels == own_labels).mean()
-    return float(purity)
-
-
-def _intra_inter_distance_ratio(distance_matrix: np.ndarray, labels: np.ndarray) -> Dict[str, float]:
-    """Mean intra-family vs inter-family pairwise distance and inter/intra ratio."""
-    n = distance_matrix.shape[0]
-    iu = np.triu_indices(n, k=1)
-    pair_d = distance_matrix[iu]
-    same = (labels[iu[0]] == labels[iu[1]])
-    intra = float(pair_d[same].mean()) if same.any() else float("nan")
-    inter = float(pair_d[~same].mean()) if (~same).any() else float("nan")
-    ratio = float(inter / intra) if intra > 0 else float("nan")
-    return {"intra_mean": intra, "inter_mean": inter, "inter_over_intra": ratio}
-
-
-def compute_structure_metrics(
-    embedding: np.ndarray,
-    labels: np.ndarray,
-    knn_k: int,
-) -> Dict[str, float]:
-    metrics: Dict[str, float] = {}
-    n_classes = len(np.unique(labels))
-    if n_classes < 2 or len(labels) < 3:
-        return {"silhouette_cosine": float("nan"),
-                "silhouette_euclidean": float("nan"),
-                "knn_purity_cosine": float("nan"),
-                "knn_purity_euclidean": float("nan"),
-                "inter_over_intra_cosine": float("nan"),
-                "inter_over_intra_euclidean": float("nan")}
-
-    metrics["silhouette_cosine"] = float(silhouette_score(embedding, labels, metric="cosine"))
-    metrics["silhouette_euclidean"] = float(silhouette_score(embedding, labels, metric="euclidean"))
-    metrics["knn_purity_cosine"] = _knn_purity(embedding, labels, knn_k, metric="cosine")
-    metrics["knn_purity_euclidean"] = _knn_purity(embedding, labels, knn_k, metric="euclidean")
-
-    cos_d = cosine_distances(embedding)
-    euc_d = euclidean_distances(embedding)
-    cos_ratio = _intra_inter_distance_ratio(cos_d, labels)
-    euc_ratio = _intra_inter_distance_ratio(euc_d, labels)
-    metrics["intra_mean_cosine"] = cos_ratio["intra_mean"]
-    metrics["inter_mean_cosine"] = cos_ratio["inter_mean"]
-    metrics["inter_over_intra_cosine"] = cos_ratio["inter_over_intra"]
-    metrics["intra_mean_euclidean"] = euc_ratio["intra_mean"]
-    metrics["inter_mean_euclidean"] = euc_ratio["inter_mean"]
-    metrics["inter_over_intra_euclidean"] = euc_ratio["inter_over_intra"]
-    return metrics
-
-
-def permutation_null(
-    embedding: np.ndarray,
-    labels: np.ndarray,
-    knn_k: int,
-    n_permutations: int,
-    seed: int,
-) -> Dict[str, Dict[str, float]]:
-    """Shuffle labels K times, recompute metrics, return per-metric (mean, std, p_value).
-
-    p-value: empirical right-tail Pr(null >= observed) under random labeling.
-    """
-    observed = compute_structure_metrics(embedding, labels, knn_k)
-    rng = np.random.default_rng(seed)
-    null_records: Dict[str, List[float]] = {k: [] for k in observed}
-
-    for _ in range(n_permutations):
-        permuted = rng.permutation(labels)
-        m = compute_structure_metrics(embedding, permuted, knn_k)
-        for k, v in m.items():
-            null_records[k].append(v)
-
-    summary: Dict[str, Dict[str, float]] = {}
-    for k, obs_v in observed.items():
-        arr = np.asarray(null_records[k], dtype=float)
-        valid = arr[~np.isnan(arr)]
-        if valid.size == 0 or np.isnan(obs_v):
-            summary[k] = {"observed": obs_v, "null_mean": float("nan"),
-                          "null_std": float("nan"), "p_right_tail": float("nan")}
-            continue
-        p = float((1 + np.sum(valid >= obs_v)) / (valid.size + 1))
-        summary[k] = {
-            "observed": float(obs_v),
-            "null_mean": float(valid.mean()),
-            "null_std": float(valid.std()),
-            "p_right_tail": p,
-        }
-    return summary
-
-
-# ---------- visualization ----------
-
-def _color_map(labels: List[str]) -> Dict[str, tuple]:
-    uniq = sorted(set(labels))
-    cmap = plt.get_cmap("tab20", max(len(uniq), 1))
-    return {lab: cmap(i) for i, lab in enumerate(uniq)}
-
-
-def _scatter(coords: np.ndarray, labels: List[str], title: str, out_path: Path) -> None:
-    fig, ax = plt.subplots(figsize=(7, 6))
-    cmap = _color_map(labels)
-    for lab in sorted(set(labels)):
-        mask = np.asarray([lab == l for l in labels])
-        ax.scatter(coords[mask, 0], coords[mask, 1],
-                   color=cmap[lab], label=f"{lab} (n={mask.sum()})",
-                   s=20, alpha=0.85, edgecolor="none")
-    ax.set_title(title)
-    ax.set_xlabel("dim 1")
-    ax.set_ylabel("dim 2")
-    ax.legend(loc="best", fontsize=8, frameon=True)
-    fig.tight_layout()
-    fig.savefig(out_path, dpi=160)
-    plt.close(fig)
-
-
-def visualize(
-    embedding: np.ndarray,
-    metadata_df: pd.DataFrame,
-    output_dir: Path,
-    umap_n_neighbors: int,
-    umap_min_dist: float,
-) -> None:
-    # PCA 2D
-    pca = PCA(n_components=2, random_state=0)
-    pca_coords = pca.fit_transform(embedding)
-    _scatter(pca_coords, metadata_df["family"].tolist(),
-             "Module embedding PCA — family", output_dir / "pca_family.png")
-
-    has_hallmark = metadata_df["hallmark_process"].notna().any()
-    if has_hallmark:
-        # Project only hallmark modules to a separate PCA for readability
-        hm_mask = metadata_df["hallmark_process"].notna().to_numpy()
-        if hm_mask.sum() >= 3:
-            hm_pca = PCA(n_components=2, random_state=0).fit_transform(embedding[hm_mask])
-            _scatter(hm_pca, metadata_df.loc[hm_mask, "hallmark_process"].tolist(),
-                     "Hallmark modules PCA — process",
-                     output_dir / "pca_hallmark_process.png")
-
-    # UMAP (optional)
-    try:
-        import umap  # type: ignore
-        reducer = umap.UMAP(
-            n_components=2,
-            n_neighbors=min(umap_n_neighbors, max(2, embedding.shape[0] - 1)),
-            min_dist=umap_min_dist,
-            metric="cosine",
-            random_state=0,
-        )
-        umap_coords = reducer.fit_transform(embedding)
-        _scatter(umap_coords, metadata_df["family"].tolist(),
-                 "Module embedding UMAP — family", output_dir / "umap_family.png")
-
-        if has_hallmark and hm_mask.sum() >= 3:
-            hm_reducer = umap.UMAP(
-                n_components=2,
-                n_neighbors=min(umap_n_neighbors, max(2, hm_mask.sum() - 1)),
-                min_dist=umap_min_dist,
-                metric="cosine",
-                random_state=0,
-            )
-            hm_umap = hm_reducer.fit_transform(embedding[hm_mask])
-            _scatter(hm_umap, metadata_df.loc[hm_mask, "hallmark_process"].tolist(),
-                     "Hallmark modules UMAP — process",
-                     output_dir / "umap_hallmark_process.png")
-    except ImportError:
-        logger.warning("umap-learn not installed; skipping UMAP plots. `pip install umap-learn` to enable.")
-
-
-# ---------- driver ----------
-
-def setup_logging(output_dir: Path) -> None:
-    log_path = output_dir / "run.log"
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
-        handlers=[logging.StreamHandler(sys.stdout), logging.FileHandler(log_path, mode="w")],
-        force=True,
-    )
-
-
-def resolve_biomod_json(args: argparse.Namespace) -> str:
-    """Resolve relative paths against biomod/cell_annotation/ to stay
-    consistent with how `load_backbone` reads `module_data_dir`.
-    """
-    if args.module_json is not None:
-        return str(ca_resolve_path(args.module_json))
-    return str(ca_resolve_path(args.module_data_dir) / "merged_biomods.json")
-
-
-def main() -> None:
-    args = parse_args()
-
-    # DDP init (auto-detects torchrun env; falls back to single-process)
-    device, use_ddp, rank, world_size, local_rank = init_distributed_mode(args.device)
-    seed_all(args.seed + rank)
-
-    if is_main_process():
-        output_dir = ensure_dir(args.output_dir)
-        setup_logging(output_dir)
-    else:
-        output_dir = Path(args.output_dir)
-        logging.basicConfig(level=logging.WARNING)
-    if use_ddp:
-        dist.barrier()
-
-    logger.info("Args: %s", vars(args))
-    logger.info("DDP: use_ddp=%s, rank=%d/%d, device=%s", use_ddp, rank, world_size, device)
-
-    # 1. Load backbone
-    backbone, biomod_data, model_config = load_backbone(args, device)
-    backbone.eval()
-    base_dim = int(model_config["base_dim"])
-    num_modules = len(biomod_data)
-    logger.info("Loaded backbone: num_modules=%d, base_dim=%d", num_modules, base_dim)
-
-    # 2. Load module names (must align with biomod_data order)
-    module_json_path = resolve_biomod_json(args)
-    module_names = load_biomod_names(module_json_path)
-    if len(module_names) != num_modules:
-        raise RuntimeError(
-            f"Module name count ({len(module_names)} from {module_json_path}) does not match "
-            f"backbone num_modules ({num_modules}). The merged_modules.json used to build the "
-            "model and the one given here must be identical, with no modules dropped during "
-            "gene-symbol resolution."
-        )
-    metadata = build_biomod_metadata(module_names)
-    metadata_df = pd.DataFrame(metadata)
-    if is_main_process():
-        logger.info("Family distribution:\n%s", metadata_df["family"].value_counts().to_string())
-
-    # 3. Build dataloader
-    if args.data_mode == "h5ad":
-        preprocessor = DataPreprocessor(
-            gene2ens_file=args.gene2ens_file,
-            gene2ens_column=args.gene2ens_column,
-            model_genes_file=args.gene_list_path,
-        )
-        preprocessor.load_gene_mapping()
-        preprocessor.load_model_genes()
-        expected_input_genes = -1
-    else:
-        preprocessor = None
-        expected_input_genes = int(model_config["num_genes"])
-
-    loader, reindex_state, n_cells = build_inference_loader(
-        args, preprocessor, use_ddp, rank, world_size,
-    )
-    if n_cells == 0:
-        raise RuntimeError("No cells survived preprocessing.")
-
-    # 4. Inference
-    logger.info("Extracting cell-averaged module embeddings over %d cells...", n_cells)
-    mean_emb, processed = extract_cell_averaged_biomod_embeddings(
-        backbone, loader, preprocessor, reindex_state, device,
-        num_modules=num_modules, base_dim=base_dim,
-        expected_input_genes=expected_input_genes,
-        use_ddp=use_ddp,
-    )
-    logger.info("Embedding ready: shape=%s, processed_cells=%d", mean_emb.shape, processed)
-
-    # --- Everything below is rank-0 only ---
-    if is_main_process():
-        # 5. Save raw outputs
-        np.save(output_dir / "module_embeddings.npy", mean_emb)
-        metadata_df.to_csv(output_dir / "module_metadata.csv", index=False)
-
-        # 6. Metrics — family level
-        if args.include_other_in_metrics:
-            keep_mask = np.ones(len(metadata_df), dtype=bool)
-        else:
-            keep_mask = (metadata_df["family"] != "Other").to_numpy()
-        family_labels = metadata_df.loc[keep_mask, "family"].tolist()
-        family_emb = mean_emb[keep_mask]
-
-        logger.info("Computing structure metrics — family level (n=%d)", keep_mask.sum())
-        family_summary = permutation_null(
-            family_emb, _label_to_int(family_labels),
-            knn_k=args.knn_k,
-            n_permutations=args.null_permutations,
-            seed=args.seed,
-        )
-
-        # 6b. Metrics — hallmark process level
-        hm_mask = metadata_df["hallmark_process"].notna().to_numpy()
-        hallmark_summary: Optional[Dict[str, Dict[str, float]]] = None
-        if hm_mask.sum() >= 3 and metadata_df.loc[hm_mask, "hallmark_process"].nunique() >= 2:
-            logger.info("Computing structure metrics — hallmark process level (n=%d)", hm_mask.sum())
-            hallmark_summary = permutation_null(
-                mean_emb[hm_mask],
-                _label_to_int(metadata_df.loc[hm_mask, "hallmark_process"].tolist()),
-                knn_k=args.knn_k,
-                n_permutations=args.null_permutations,
-                seed=args.seed,
-            )
-
-        metrics_payload = {
-            "n_cells_processed": int(processed),
-            "num_modules": int(num_modules),
-            "base_dim": int(base_dim),
-            "knn_k": int(args.knn_k),
-            "null_permutations": int(args.null_permutations),
-            "family_metrics": family_summary,
-            "hallmark_process_metrics": hallmark_summary,
-            "family_counts": metadata_df["family"].value_counts().to_dict(),
-            "hallmark_process_counts": metadata_df["hallmark_process"].dropna().value_counts().to_dict(),
-            "shuffle_module_genes": bool(args.shuffle_module_genes),
-        }
-        with open(output_dir / "metrics.json", "w", encoding="utf-8") as fh:
-            json.dump(metrics_payload, fh, indent=2)
-
-        logger.info("=== Family-level structure (vs random-shuffled label null) ===")
-        for k, v in family_summary.items():
-            logger.info("  %-30s observed=%.4f  null=%.4f±%.4f  p=%.4f",
-                        k, v["observed"], v["null_mean"], v["null_std"], v["p_right_tail"])
-        if hallmark_summary is not None:
-            logger.info("=== Hallmark-process-level structure (Hallmarks subset) ===")
-            for k, v in hallmark_summary.items():
-                logger.info("  %-30s observed=%.4f  null=%.4f±%.4f  p=%.4f",
-                            k, v["observed"], v["null_mean"], v["null_std"], v["p_right_tail"])
-
-        # 7. Visualisation
-        logger.info("Plotting UMAP / PCA...")
-        visualize(mean_emb, metadata_df, output_dir,
-                  umap_n_neighbors=args.umap_n_neighbors,
-                  umap_min_dist=args.umap_min_dist)
-
-        logger.info("Done. Outputs saved to %s", output_dir)
-
-    # Teardown
-    if use_ddp:
-        dist.barrier()
-        dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/README.md

@@ -1,319 +0,0 @@
-# Experiment 3: Module Importance and Enrichment Analysis
-
-This experiment tests whether modules that a finetuned model considers
-important are consistent with known cell-type marker genes.
-
-## Method
-
-1. Load a finetuned cell annotation model: backbone, attention-pooling head, and classifier.
-2. Mask one module at a time on the test set and measure the drop in the correct-class logit (`DeltaLogit`).
-3. Aggregate by cell type to obtain an importance matrix with shape `(num_classes x 660)`.
-4. For the top-k modules of each cell type, run hypergeometric enrichment on their genes.
-
-## Masking Modes
-
-| Mode | Principle | Speed |
-| --- | --- | --- |
-| `post_encoder` | Zero out the module token after encoder output | Fast, because the encoder runs once |
-| `pre_encoder` | Zero out the genes corresponding to the module before encoder input | Slow, because it requires 660 x cells forward passes |
-
-`pre_encoder` is more accurate because it captures interactions between modules,
-but it is much slower. Use `--importance_cells` for subsampling control.
-
-## Run
-
-Run these commands from the BioMod repository root. First train a matching
-cell-annotation model; the example checkpoint path is the output of
-`BioMod/toturials/cell_annotation.ipynb`, not a bundled pretrained weight file.
-Use the matching BioMod configuration, classifier architecture, preprocessing,
-label mapping, and test split for that trained model. Substitute your own
-`--finetune_checkpoint` if its output directory differs. Importance matrices,
-ablation outputs, and permutation results below must be generated by the
-preceding steps; they are not bundled inputs.
-
-```bash
-# Run from the BioMod repository root.
-python BioMod/interpretability/biomod_importance/biomod_importance.py \
-    --finetune_checkpoint tutorial_runs/cell_annotation_zheng68k_biomod/best_model.pt \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --mask_mode post_encoder \
-    --batch_size 128 \
-    --top_k 10 20 \
-    --output_dir outputs/interpretability/importance/zheng68k_post
-```
-
-Pre-encoder mode, more accurate but slower:
-
-```bash
-python BioMod/interpretability/biomod_importance/biomod_importance.py \
-    --finetune_checkpoint tutorial_runs/cell_annotation_zheng68k_biomod/best_model.pt \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --mask_mode pre_encoder \
-    --importance_cells 500 \
-    --batch_size 256 \
-    --top_k 10 20 \
-    --output_dir outputs/interpretability/importance/zheng68k_pre
-```
-
-Pre-encoder mode on all cells with DDP multi-GPU acceleration:
-
-```bash
-torchrun --nproc_per_node=4 BioMod/interpretability/biomod_importance/biomod_importance.py \
-    --finetune_checkpoint tutorial_runs/cell_annotation_zheng68k_biomod/best_model.pt \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --mask_mode pre_encoder \
-    --batch_size 512 \
-    --top_k 10 20 \
-    --output_dir outputs/interpretability/importance/zheng68k_pre_full
-```
-
-## Arguments
-
-| Argument | Default | Description |
-| --- | --- | --- |
-| `--mask_mode` | post_encoder | Masking strategy |
-| `--importance_cells` | 0 | Maximum number of cells for `pre_encoder`; 0 means all cells |
-| `--top_k` | 10 20 | Top-k modules used for enrichment analysis |
-| `--batch_size` | 64 | Inference batch size |
-| `--test_ratio` | 0.2 | Test-set fraction; should match finetuning |
-| `--seed` | 2021 | Random seed; should match finetuning |
-| `--marker_db_dir` | None | Marker database cache directory |
-
-## Outputs
-
-```text
-output_dir/
-|-- importance_matrix.npy           # (11, 660) DeltaLogit matrix
-|-- importance_matrix.csv           # Same matrix with labels
-|-- top_modules_per_celltype.json   # Top modules and genes for each cell type
-|-- enrichment_results.csv          # Hypergeometric enrichment results
-|-- heatmap_importance.png          # Heatmap
-|-- top_modules_bar.png             # Bar plot of top modules per cell type
-|-- metrics.json                    # Summary metrics
-`-- marker_cache/                   # PanglaoDB cache
-    `-- PanglaoDB_markers.tsv
-```
-
-## Expected Conclusion
-
-If the model learns meaningful module importance:
-
-- Top modules for immune cell types should be enriched for immune-related pathways.
-- Enrichment p-values should be significant after Bonferroni correction, below 0.05.
-- Marker genes for the corresponding cell types in PanglaoDB should appear in top modules.
-
----
-
-## Extended Experiments
-
-### 3-A: Masking Ablation Curves
-
-Progressively mask top, bottom, or random modules and measure the accuracy drop.
-
-```bash
-python BioMod/interpretability/biomod_importance/masking_ablation.py \
-    --finetune_checkpoint tutorial_runs/cell_annotation_zheng68k_biomod/best_model.pt \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --importance_matrix outputs/interpretability/importance/zheng68k_pre_13690/importance_matrix.npy \
-    --k_values 1 2 5 10 20 50 100 200 \
-    --n_random_trials 10 \
-    --batch_size 360 \
-    --save_deltas \
-    --output_dir outputs/interpretability/importance/ablation/
-```
-
-### 3-B: Permutation Tests
-
-```bash
-# Mode A: shuffle modules, GPU and slow.
-python BioMod/interpretability/biomod_importance/random_baseline.py \
-    --mode shuffle_modules \
-    --finetune_checkpoint tutorial_runs/cell_annotation_zheng68k_biomod/best_model.pt \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --importance_matrix outputs/interpretability/importance/zheng68k_pre_13690/importance_matrix.npy\
-    --n_permutations 10 \
-    --output_dir outputs/interpretability/importance/permutation/
-
-# Mode B: shuffle labels, CPU and fast, requires per_cell_deltas from ablation.
-python BioMod/interpretability/biomod_importance/random_baseline.py \
-    --mode shuffle_labels \
-    --importance_matrix outputs/interpretability/importance/zheng68k_pre_13690/importance_matrix.npy \
-    --per_cell_deltas outputs/interpretability/importance/ablation/per_cell_deltas.npy \
-    --per_cell_labels outputs/interpretability/importance/ablation/per_cell_labels.npy \
-    --n_permutations 1000 \
-    --output_dir outputs/interpretability/importance/permutation/
-```
-
-### 3-C: Marker Gene Verification, CPU Only
-
-```bash
-python BioMod/interpretability/biomod_importance/marker_gene_verification.py \
-    --importance_matrix outputs/interpretability/importance/zheng68k_pre_13690/importance_matrix.npy \
-    --importance_csv <importance_matrix.csv> \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --output_dir outputs/interpretability/importance/marker_verification/
-```
-
-### Publication Figures
-
-```bash
-python BioMod/interpretability/biomod_importance/biomod_importance_publication.py \
-    --results_dir outputs/interpretability/importance/zheng68k_pre_13690/ \
-    --ablation_dir outputs/interpretability/importance/ablation \
-    --permutation_dir outputs/interpretability/importance/permutation/ \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --output_dir outputs/interpretability/importance/publication_figures/
-```
-
-### 3-E: Family-Level Analysis, CPU Only
-
-Analyze importance distributions, contribution fractions, and enrichment tests by
-module source family: Hallmarks, Reactome, Hotspot, and WGCNA.
-
-```bash
-python BioMod/interpretability/biomod_importance/family_importance_analysis.py \
-    --results_dir outputs/interpretability/importance/zheng68k_pre \
-    --output_dir outputs/interpretability/importance/family_analysis \
-    --top_k 10 20 50 \
-    --enrichment_k 50
-```
-
-Outputs:
-
-- `fig1_family_importance_distribution.png/pdf`: family-wise importance distribution for each cell type, violin + box.
-- `fig2_family_contribution_topk.png/pdf`: family fraction among top-K important modules, stacked bar.
-- `fig3_family_enrichment_heatmap.png/pdf`: Fisher exact-test enrichment-ratio heatmap.
-- `fig4_family_mean_importance.png/pdf`: family mean-importance comparison, grouped bar.
-- `kruskal_wallis_family_test.csv`: Kruskal-Wallis test results.
-- `pairwise_mannwhitney_families.csv`: pairwise Mann-Whitney U tests between families.
-- `family_enrichment_fisher.csv`: detailed Fisher exact-test results with Bonferroni correction.
-- `family_contribution_topk.csv`: numerical family contribution data.
-- `family_mean_importance.csv`: family mean-importance matrix.
-
-### 3-F: Family-Level Ablation, GPU
-
-Ablate whole families by removing or keeping one module family and measuring
-classification-accuracy changes.
-
-```bash
-python BioMod/interpretability/biomod_importance/family_ablation.py \
-    --finetune_checkpoint tutorial_runs/cell_annotation_zheng68k_biomod/best_model.pt \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --batch_size 256 \
-    --output_dir outputs/interpretability/importance/family_analysis_gpu
-```
-
-Outputs:
-
-- `family_ablation_results.json`: overall and per-class accuracy under all conditions.
-- `fig_family_ablation.png/pdf`: two-panel overall plot for remove-one vs keep-one.
-- `fig_family_ablation_per_class.png/pdf`: two-panel heatmap of per-cell-type ablation results.
-
-### 3-G: Cross-Class Similarity, CPU Only
-
-Analyze correlations and similarity between important modules across the 11 cell
-classes.
-
-```bash
-python BioMod/interpretability/biomod_importance/cross_class_similarity.py \
-    --results_dir outputs/interpretability/importance/zheng68k_pre \
-    --output_dir outputs/interpretability/importance/publication_figures \
-    --top_k 50
-```
-
-### 3-H: 1B Module Gating + 2B Family Embedding Finetuning Experiment
-
-Based on module-importance findings, most modules have near-zero importance and
-different cell types depend on different families. During finetuning, add two
-lightweight modules in a LoRA-style setup: pretrained weights are frozen and new
-parameters are trained from zero initialization.
-
-- **1B, Sparse Module Gate**: `LayerNorm -> Linear(E,1) -> Sigmoid`, learning a 0-1 gate scalar for each module.
-- **2B, Family Embedding**: `nn.Parameter(4, E)`, one learnable vector per family, added to `pos_emb`.
-
-```bash
-# Baseline, no new modules.
-python BioMod/cell_annotation/biomod_finetune.py \
-    --config_pkl weight/BioMod.pkl \
-    --biomod_checkpoint weight/BioMod.pth \
-    --biomod_data_dir model_dependencies/biomod_data_all \
-    --gene_list_path model_dependencies/selected_genes_27k.txt \
-    --gene_mapping_path model_dependencies/gene2ens.csv \
-    --gene2vec_data_path model_dependencies/selected_gene2vec_27k.npy \
-    --gene2ens_file model_dependencies/gene2ens.csv \
-    --data_path data/downstream/final_data/Zheng68K.h5ad \
-    --label_column celltype \
-    --head_type attn_pool \
-    --finetune_mode last_n --unfreeze_last_n 1 \
-    --lr_head 0.0001 --lr_backbone 0.0001 \
-    --batch_size 256 --epochs 300 \
-    --run_name baseline_zheng68k \
-    --output_dir outputs/interpretability/family_exp
-```
-
-The same command can be extended with:
-
-- `--use_family_emb` to enable family-type embedding, adding `4 x 200 = 800` parameters.
-- `--use_biomod_gate` to enable sparse module gating, adding LayerNorm + Linear, 601 parameters.
-- Both flags can be used together.
-
-The new module parameters use `lr_head`, a higher learning rate, and are
-independent of frozen backbone parameters.
-
-## File Descriptions
-
-| File | Description |
-| --- | --- |
-| `biomod_importance.py` | Core importance computation with DDP support |
-| `masking_ablation.py` | Ablation-curve experiment, GPU |
-| `random_baseline.py` | Permutation tests, GPU/CPU |
-| `marker_gene_verification.py` | Marker verification, CPU |
-| `biomod_importance_publication.py` | Publication figures |
-| `family_importance_analysis.py` | Family-level importance analysis, CPU |
-| `family_ablation.py` | Family-level ablation, GPU |
-| `cross_class_similarity.py` | Cross-class similarity analysis, CPU |
-| `pub_utils.py` | Shared utilities |
-| `RESULTS.md` | Result documentation with figure placeholders |
-| `TODO.md` | Progress tracker |

BioMod/interpretability/biomod_importance/RESULTS.md

@@ -1,226 +0,0 @@
-# Module Importance: Result Documentation
-
-> For the manuscript subsection "Model Interpretability: Module Importance Analysis".
-
-## Dataset
-
-- **Zheng68K**: 11 cell types, 68,450 cells, with 13,690 test cells.
-- Model: BioModCellClassifier with attention-pooling head and 48M backbone.
-- Baseline test accuracy: **83.1%**.
-
-## Data Sources
-
-These are historical experiment results. The analysis output files below are not
-included in the repository; recreate them using the commands in `README.md` or
-point the analysis scripts to the original run directory. The shared module JSON
-is bundled under `model_dependencies/biomod_data_all/` at the repository root.
-
-All analyses are based on the following core files, generated by `biomod_importance.py`
-with `pre_encoder` over all 13,690 test cells:
-
-| File | Path | Description |
-| --- | --- | --- |
-| `importance_matrix.npy` | `outputs/interpretability/importance/zheng68k_pre_13690/importance_matrix.npy` | `(11 x 660)` DeltaLogit matrix |
-| `importance_matrix.csv` | `outputs/interpretability/importance/zheng68k_pre_13690/importance_matrix.csv` | Same matrix with row and column labels |
-| `enrichment_results.csv` | `outputs/interpretability/importance/zheng68k_pre_13690/enrichment_results.csv` | Enrichment-test results |
-| `top_modules_per_celltype.json` | `outputs/interpretability/importance/zheng68k_pre_13690/top_modules_per_celltype.json` | Top-module details |
-| `merged_biomods.json` | `model_dependencies/biomod_data_all/merged_biomods.json` | Gene composition of 660 modules |
-
-Marker verification uses `importance_matrix.npy`, `importance_matrix.csv`, and
-`merged_biomods.json`.
-
-Publication figures use all files above plus ablation and permutation results.
-
----
-
-## Experiment 3 Core: Module Importance Matrix
-
-### Method
-
-For each test cell, mask one module at a time using `pre_encoder`, which zeros out
-the corresponding gene inputs, and measure the drop in the correct-class logit
-as the importance score. Aggregate by cell type to obtain an `(11 x 660)`
-importance matrix.
-
-### Results
-
-- Masking mode: **pre_encoder** over all 13,690 cells.
-- 660 modules: 491 Reactome, 50 Hallmarks, 73 Hotspot, and 46 Processed.
-
-**Representative top modules per cell type:**
-
-| Cell Type | #1 Module | #2 Module | #3 Module |
-| --- | --- | --- | --- |
-| CD14+ Monocyte | Neutrophil Degranulation | TLR4 Cascade | MHC-II Antigen Presentation |
-| CD19+ B | MHC-II | IFN-gamma Signaling | CD28 Co-stimulation |
-| CD56+ NK | SM_22 (cytotoxic) | module3_2 | Allograft Rejection |
-| CD8+ Cytotoxic T | module3_2 | SM_22 | Allograft Rejection |
-| Dendritic | MHC-II | Signaling by Interleukins | IFN-gamma |
-
-### Figure
-
-Expected generated figure: `publication_figures/fig4_importance_heatmap_clustered.png` (Importance Heatmap (clustered); not bundled).
-
----
-
-## Experiment 3-A: Masking Ablation Curves
-
-### Method
-
-Progressively mask top-k, bottom-k, or random-k modules, with
-`k = 1, 2, 5, 10, 20, 50, 100, 200`, and measure classification accuracy. This
-uses `post_encoder` masking so the encoder runs only once.
-
-### Expected Result
-
-- Masking top modules should rapidly decrease accuracy.
-- Masking bottom modules should have little effect.
-- Random masking should fall between the two.
-
-### Results, Completed
-
-- Baseline accuracy: **83.0%** on 13,690 test cells.
-- Post-encoder masking with k = [1, 2, 5, 10, 20, 50, 100, 200].
-
-| Masked k | Top Acc | Bottom Acc | Random Acc |
-| --- | --- | --- | --- |
-| 1 | 83.0% | 83.0% | 83.0% |
-| 2 | 82.6% | 83.0% | 83.0% |
-| 5 | 82.0% | 83.0% | 83.0% |
-| 10 | 77.1% | 83.0% | 83.0% |
-| 20 | 73.7% | 83.1% | 82.7% |
-| 50 | 69.5% | 83.2% | 82.4% |
-| 100 | 65.1% | 82.7% | 81.7% |
-| 200 | 49.7% | 81.9% | 78.6% |
-
-**Interpretation:** masking the top 200 modules, 30% of all modules, reduces
-accuracy from 83% to 50%, a 33 percentage-point drop. Masking the bottom 200
-modules has almost no effect, with only a 1.1 point drop. Random masking falls
-between the two. This shows that the importance ranking captures modules that
-are truly important for classification.
-
-Expected generated figure: `publication_figures/fig1_ablation_curves.png` (Ablation Curves; not bundled).
-
-Expected generated figure: `publication_figures/fig_s1_ablation_per_class.png` (Per-class Ablation; not bundled).
-
----
-
-## Experiment 3-B: Permutation Tests, Random Baseline
-
-### Method
-
-**Mode A, Module-Shuffle:** shuffle module token order to break the module-gene
-correspondence, then recompute importance. Repeat 10 times.
-
-**Mode B, Label-Shuffle:** shuffle cell type labels and re-aggregate per-cell
-deltas. Repeat 1000 times.
-
-### Expected Result
-
-- The true importance maximum should far exceed the null distribution.
-- Empirical p-value should be below 0.001.
-
-### Results, Completed
-
-**Mode A, Module-Shuffle** (10 permutations, GPU):
-
-| Metric | Value |
-| --- | --- |
-| Global p-value | 0.091, limited by n=10; minimum possible is 1/11 |
-| Effect size | **1474.5** |
-| Real global max | 2.046 |
-| Null global max | 0.452 +/- 0.001 |
-
-**Mode B, Label-Shuffle** (1000 permutations, CPU):
-
-| Metric | Value |
-| --- | --- |
-| Global p-value | **< 0.001** |
-| Per-class p-values | All < 0.001 (11/11) |
-| Effect size | **48.1** |
-| Real global max | 2.046 |
-| Null global max | 0.145 +/- 0.039 |
-
-**Interpretation:**
-
-- **Module-Shuffle:** after shuffling the module-gene correspondence, the maximum importance drops from 2.046 to about 0.452. The effect size is extremely large, showing that importance depends on the correct module-gene mapping rather than random noise. The p-value is 0.091 because only 10 permutations were run, making the minimum possible p-value about 1/11; in practice, 0 of 10 null runs exceeded the observed value.
-- **Label-Shuffle:** after shuffling cell type labels, the maximum importance drops from 2.046 to about 0.145. All 11 cell types have p < 0.001, showing that importance is cell-type-specific rather than a generic pattern shared by all cells.
-
-Expected generated figure: `publication_figures/fig2_permutation_null.png` (Permutation Null; not bundled).
-
----
-
-## Experiment 3-C: Canonical Marker Gene Verification
-
-### Method
-
-Check the importance ranks of modules containing known cell-type canonical
-markers such as CD14, CD19, CD8A, and NKG7. Test whether markers are enriched in
-top modules.
-
-### Results, Completed
-
-| Metric | Value |
-| --- | --- |
-| Total markers tested | 59 |
-| Markers not found | 2 |
-| **Median rank percentile** | **0.9%** |
-| Markers in top 5% | **47/59 (80%)** |
-| Markers in top 10% | 47/59 (80%) |
-| Markers in top 20% | 50/59 (85%) |
-| Binomial p-value, top 10% | **< 1e-50** |
-| Wilcoxon p-value, <50% | **2.5e-11** |
-
-**Per cell type:**
-
-| Cell Type | Median Percentile | Best Marker | Best Rank |
-| --- | --- | --- | --- |
-| CD14+ Monocyte | 0.0% | LYZ | #1 |
-| CD19+ B | 1.3% | MS4A1 | #7 |
-| CD34+ | 18.5% | THY1 | #7 |
-| CD4+ T Helper2 | 1.7% | CD4 | #1 |
-| CD4+/CD25 T Reg | 0.9% | CD4 | #1 |
-| CD4+/CD45RA+/CD25- Naive T | 10.5% | CD4 | #1 |
-| CD4+/CD45RO+ Memory | 0.0% | CD4 | #1 |
-| CD56+ NK | 0.0% | NKG7 | #1 |
-| CD8+ Cytotoxic T | 0.0% | CD8A | #1 |
-| CD8+/CD45RA+ Naive Cytotoxic | 20.3% | CD8A | #1 |
-| Dendritic | 2.1% | ITGAX | #1 |
-
-**Interpretation:** 80% of canonical markers are in the top 5% importance
-modules for the corresponding cell type. In 9 of 11 cell types, the best marker
-is in the rank #1 module. This strongly shows that the model learns biologically
-meaningful module importance.
-
-Expected generated figure: `marker_verification/fig_marker_ranks.png` (Marker Ranks; not bundled).
-
-Expected generated figure: `marker_verification/fig_marker_summary_bar.png` (Marker Summary Bar; not bundled).
-
----
-
-## Experiment 3-D: PanglaoDB Enrichment Analysis
-
-### Method
-
-For genes in each cell type's top-k modules, run a hypergeometric enrichment test
-against marker gene sets for 178 PanglaoDB cell types, with Bonferroni
-correction.
-
-### Results
-
-| Metric | Value |
-| --- | --- |
-| Significant enrichment hits, p_adj < 0.05 | **1011 / 3916** |
-| Hit rate | **25.8%** |
-
-Expected generated figure: `publication_figures/fig5_enrichment_summary.png` (Enrichment Summary; not bundled).
-
----
-
-## Manuscript Writing Points
-
-1. **Core conclusion:** the model learns biologically meaningful gene organization through the module-based architecture. Top-importance modules strongly match known cell-type markers, with 80% of markers in the top 5%, and masking these modules causes a large classification-performance drop, from 83% to 50%.
-2. **Methodological contribution:** pre-encoder masking gives a more accurate importance estimate than post-encoder masking because it accounts for interactions between modules. DDP support makes all-cell computation feasible.
-3. **Statistical rigor:** label-shuffle p < 0.001 with effect size 48.1 shows that importance is cell-type-specific; module-shuffle effect size 1474 shows that importance depends on the correct module-gene correspondence. Top-k masking rapidly reduces accuracy while bottom-k masking has almost no effect, and the gap widens as k increases.
-4. **Biological interpretation:** monocytes use Neutrophil Degranulation and TLR4, representing innate immune pathways; B cells use MHC-II, IFN-gamma, and CD28, representing antigen presentation and T-B cooperation; NK/CD8 T share cytotoxic modules; dendritic cells use MHC-II and interleukins, matching professional antigen-presenting cells.
-5. **Key numbers for the main text:** 83.0% baseline accuracy; 80% canonical markers in top 5% importance modules; 9 of 11 cell types have the best marker in the rank #1 module; masking top-200 modules drops accuracy by 33 percentage points; label-shuffle permutation p < 0.001 with effect size 48.1; PanglaoDB enrichment hit rate 25.8% (1011/3916 significant).

BioMod/interpretability/biomod_importance/TODO.md

@@ -1,68 +0,0 @@
-# Module Importance: TODO / Progress Tracker
-
-## Status Legend
-
-- [x] Completed
-- [ ] To do
-- [~] In progress
-
----
-
-## Phase 1: CPU-Only Analysis, No GPU Required
-
-- [x] `pub_utils.py`: shared utilities for colors, styles, markers, and loading functions.
-- [x] `marker_gene_verification.py`: canonical marker verification.
-  - Result: 80% of markers are in the top 5%, median percentile = 0.9%.
-  - Figures: `fig_marker_ranks.pdf`, `fig_marker_summary_bar.pdf`.
-- [x] `biomod_importance_publication.py`: publication-figure generator.
-- [x] `RESULTS.md`: result documentation.
-- [x] `TODO.md`: this file.
-
----
-
-## Phase 2: GPU Experiments, Run on Server
-
-- [x] `masking_ablation.py`: masking ablation curves.
-  - Result: top-200 masking reduces accuracy from 83% to 50%, while bottom-200 has almost no effect.
-  - Outputs: `ablation_curves.npz`, `ablation_metrics.json`, `per_cell_deltas.npy`, `attention_weights.npy`.
-- [x] `random_baseline.py` Mode A: module-shuffle permutation, GPU.
-  - Result: effect size = 1474, null max = 0.452 vs real max = 2.046.
-  - Output: `null_module_shuffle.npy`.
-- [x] `random_baseline.py` Mode B: label-shuffle permutation, CPU.
-  - Result: p < 0.001 for all 11 classes, effect size = 48.1.
-  - Output: `null_label_shuffle.npy`.
-
----
-
-## Phase 3: Publication Figures
-
-- [x] `fig1_ablation_curves`: ablation curves for top, bottom, and random modules.
-- [x] `fig2_permutation_null`: permutation-test null distribution.
-- [x] `fig4_importance_heatmap_clustered`: clustered heatmap.
-- [x] `fig5_enrichment_summary`: PanglaoDB enrichment summary.
-- [x] `fig_s1_ablation_per_class`: ablation curves per cell type.
-- [x] `fig_marker_ranks`: marker gene rank percentile.
-- [x] `fig_marker_summary_bar`: marker summary bar plot.
-
----
-
-## Phase 4: Extensions, Optional
-
-- [ ] Cross-dataset validation, requiring finetuned checkpoints on other datasets.
-- [ ] Correlation analysis between attention weights and masking importance.
-- [ ] Module specificity analysis, using entropy across cell types.
-- [ ] Secondary GO/KEGG enrichment analysis.
-
----
-
-## Completion Summary
-
-The core experiments are complete and the conclusion is clear:
-
-1. **Ablation** shows that top modules are necessary for classification, because accuracy drops strongly after masking them.
-2. **Permutation** shows that importance is non-random under module-shuffle and cell-type-specific under label-shuffle.
-3. **Marker verification** shows that the model learns biologically meaningful module importance, with 80% of markers in the top 5%.
-4. **Enrichment** shows that top modules are enriched for marker genes of the corresponding cell type, with a 25.8% hit rate.
-
-These results are sufficient to support a complete Model Interpretability
-subsection in the manuscript.

BioMod/interpretability/biomod_importance/biomod_importance.py

@@ -1,851 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Experiment ③ — Module Importance + Enrichment Analysis
-
-Measures per-cell-type module importance via single-module masking on a
-finetuned cell annotation model, then validates top modules against
-known cell-type marker gene databases (PanglaoDB / CellMarker).
-
-Two masking modes:
-- pre_encoder:  zero-out module *input* before encoder (accurate, slow)
-- post_encoder: zero-out module *token* after encoder (fast, encoder runs once)
-"""
-
-import argparse
-import json
-import logging
-import os
-import pickle
-import sys
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-from scipy import sparse, stats
-from sklearn.model_selection import StratifiedShuffleSplit
-from sklearn.preprocessing import LabelEncoder
-from torch.utils.data import DataLoader, Dataset, Subset
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-INTERP_DIR = SCRIPT_DIR.parent
-BIOMOD_ROOT = INTERP_DIR.parent
-SCLONG_ROOT = BIOMOD_ROOT.parent
-
-for p in (str(BIOMOD_ROOT), str(SCLONG_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger(__name__)
-
-
-# ─── DDP Utilities ─────────────────────────────────────────────────────────
-
-def setup_ddp():
-    """Initialize DDP if launched via torchrun, otherwise return single-GPU info."""
-    if "RANK" in os.environ:
-        dist.init_process_group(backend="nccl")
-        rank = dist.get_rank()
-        world_size = dist.get_world_size()
-        local_rank = int(os.environ.get("LOCAL_RANK", 0))
-        torch.cuda.set_device(local_rank)
-        device = torch.device(f"cuda:{local_rank}")
-        if rank != 0:
-            logging.getLogger().setLevel(logging.WARNING)
-        return rank, world_size, device, True
-    else:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        return 0, 1, device, False
-
-
-def cleanup_ddp(use_ddp: bool):
-    if use_ddp:
-        dist.destroy_process_group()
-
-
-# ─── Dataset ────────────────────────────────────────────────────────────────
-
-class CellDataset(Dataset):
-    def __init__(self, X, labels):
-        self.X = X
-        self.labels = labels
-
-    def __len__(self):
-        return self.X.shape[0]
-
-    def __getitem__(self, idx):
-        row = self.X[idx]
-        if sparse.issparse(row):
-            row = row.toarray().squeeze(0)
-        return torch.tensor(row, dtype=torch.float32), self.labels[idx]
-
-
-# ─── Model Loading ──────────────────────────────────────────────────────────
-
-def load_finetuned_model(args, device: torch.device):
-    """Load the full finetuned classifier (backbone + pooler + head)."""
-    from models_biomod import BioMod
-    from biomod_dataset import load_biomod_data
-
-    ckpt_path = Path(args.finetune_checkpoint)
-    checkpoint = torch.load(ckpt_path, map_location=device)
-
-    saved_args = checkpoint.get("args", {})
-    backbone_config = checkpoint.get("backbone_config", {})
-    label_classes = checkpoint["label_classes"]
-    num_classes = len(label_classes)
-
-    config_pkl_path = Path(args.config_pkl)
-    config_pkl = None
-    if config_pkl_path.exists():
-        with open(config_pkl_path, "rb") as f:
-            config_pkl = pickle.load(f)
-
-    biomod_data = None
-    if isinstance(config_pkl, dict) and "module_sets" in config_pkl:
-        biomod_data = config_pkl["module_sets"]
-    elif "module_sets" in backbone_config:
-        biomod_data = backbone_config["module_sets"]
-    else:
-        biomod_data = load_biomod_data(
-            module_data_dir=args.module_data_dir,
-            gene_list_path=args.gene_list_path,
-            gene_mapping_path=args.gene_mapping_path,
-            verbose=True,
-        )
-
-    base_dim = backbone_config.get("base_dim", 1024)
-    enc_depth = backbone_config.get("enc_depth", 32)
-    enc_heads = backbone_config.get("enc_heads", 16)
-    enc_dim_head = backbone_config.get("enc_dim_head", 64)
-    num_genes = backbone_config.get("num_genes", 27874)
-    use_cls_token = backbone_config.get("use_cls_token", False)
-    gene_embed_dim = backbone_config.get("gene_embed_dim", 200)
-    fusion_method = backbone_config.get("fusion_method", "add")
-    use_gene2vec = backbone_config.get("use_gene2vec", True)
-    loss_mode = backbone_config.get("loss_mode", "module_mean")
-    recon_target_level = backbone_config.get("recon_target_level", "module")
-    pretrain_arch = backbone_config.get("pretrain_arch", "mae")
-    mask_style = backbone_config.get("mask_style", "random")
-    mask_token_mode = backbone_config.get("mask_token_mode", "learnable")
-    pos_embed_mode = backbone_config.get("pos_embed_mode", "sincos")
-    decoder_type = backbone_config.get("decoder_type", "cross_attn")
-    decoder_gene_agg = backbone_config.get("decoder_gene_agg", "mean")
-    dec_embed_dim = backbone_config.get("dec_embed_dim", 512)
-    dec_depth = backbone_config.get("dec_depth", 4)
-    dec_heads = backbone_config.get("dec_heads", 8)
-    dec_dim_head = backbone_config.get("dec_dim_head", 64)
-
-    backbone = BioMod(
-        module_sets=biomod_data,
-        num_genes=num_genes,
-        max_seq_len=len(biomod_data),
-        base_dim=base_dim,
-        enc_depth=enc_depth,
-        enc_heads=enc_heads,
-        enc_dim_head=enc_dim_head,
-        mask_ratio=0.75,
-        gene2vec_file=args.gene2vec_data_path,
-        gene_embed_dim=gene_embed_dim,
-        fusion_method=fusion_method,
-        use_cls_token=use_cls_token,
-        cls_loss_weight=0.0,
-        cls_head_hidden_dim=0,
-        use_gene2vec=use_gene2vec,
-        loss_mode=loss_mode,
-        recon_target_level=recon_target_level,
-        pretrain_arch=pretrain_arch,
-        mask_style=mask_style,
-        mask_token_mode=mask_token_mode,
-        pos_embed_mode=pos_embed_mode,
-        decoder_type=decoder_type,
-        decoder_gene_agg=decoder_gene_agg,
-        dec_embed_dim=dec_embed_dim,
-        dec_depth=dec_depth,
-        dec_heads=dec_heads,
-        dec_dim_head=dec_dim_head,
-    )
-
-    sys.path.insert(0, str(BIOMOD_ROOT / "cell_annotation"))
-    from biomod_finetune import (
-        GatedAttentionPool,
-        BioModCellClassifier,
-        ScGPTClsDecoder,
-    )
-
-    head_type = saved_args.get("head_type", args.head_type)
-    dropout = saved_args.get("dropout", 0.2)
-    pool_hidden_dim = saved_args.get("pool_hidden_dim", 0)
-    cls_decoder_layers = saved_args.get("cls_decoder_layers", 3)
-
-    model = BioModCellClassifier(
-        backbone=backbone,
-        module_sets=biomod_data,
-        num_classes=num_classes,
-        head_type=head_type,
-        pool_hidden_dim=pool_hidden_dim,
-        dropout=dropout,
-        cls_decoder_layers=cls_decoder_layers,
-    )
-
-    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
-    model.to(device)
-    model.eval()
-
-    logger.info(
-        "Loaded finetuned model: head=%s, classes=%d, modules=%d",
-        head_type, num_classes, len(biomod_data),
-    )
-    return model, biomod_data, label_classes, head_type
-
-
-# ─── Data Loading ───────────────────────────────────────────────────────────
-
-def load_test_data(args, label_classes: List[str], device: torch.device):
-    """Load Zheng68K and reproduce the test split."""
-    import anndata as ad
-
-    sys.path.insert(0, str(BIOMOD_ROOT / "cell_annotation"))
-    from biomod_finetune import DataPreprocessor
-
-    preprocessor = DataPreprocessor(
-        gene2ens_file=args.gene2ens_file,
-        gene2ens_column="Ensembl gene ID",
-        model_genes_file=args.gene_list_path,
-    )
-    preprocessor.load_gene_mapping()
-    preprocessor.load_model_genes()
-
-    adata, reindex_state = preprocessor.preprocess_adata(
-        adata_path=args.data_path,
-        min_genes=10,
-        min_cells=10,
-        dataset_name="test_data",
-        gene_name_column=args.gene_name_column,
-    )
-
-    labels_raw = adata.obs[args.label_column].values
-    le = LabelEncoder()
-    le.classes_ = np.array(label_classes)
-    labels_encoded = le.transform(labels_raw)
-
-    indices = np.arange(len(labels_encoded))
-    train_val_split = StratifiedShuffleSplit(
-        n_splits=1, test_size=args.test_ratio, random_state=args.seed
-    )
-    train_val_idx, test_idx = next(train_val_split.split(indices, labels_encoded))
-
-    X_test = adata.X[test_idx]
-    y_test = labels_encoded[test_idx]
-
-    logger.info("Test set: %d cells, %d classes", len(test_idx), len(le.classes_))
-    return X_test, y_test, le, preprocessor, reindex_state
-
-
-# ─── Importance Computation ─────────────────────────────────────────────────
-
-@torch.no_grad()
-def compute_importance_post_encoder(
-    model: nn.Module,
-    dataloader: DataLoader,
-    preprocessor,
-    reindex_state: dict,
-    num_modules: int,
-    num_classes: int,
-    device: torch.device,
-) -> np.ndarray:
-    """Fast mode: encoder runs once, mask each module token after encoder."""
-    importance = np.zeros((num_classes, num_modules), dtype=np.float64)
-    class_counts = np.zeros(num_classes, dtype=np.float64)
-
-    for x_batch, y_batch in dataloader:
-        x_batch = x_batch.to(device)
-        y_batch = y_batch.numpy()
-        x_model = preprocessor.reindex_tensor_to_model_space(
-            x_batch, reindex_state=reindex_state
-        )
-
-        latent, _, ids_restore, _ = model.backbone.forward_encoder(
-            x_model, mask_ratio=0.0
-        )
-        if model.head_type == "cls":
-            module_tokens = latent[:, 1:, :]
-        else:
-            module_tokens = model._restore_biomod_order(latent, ids_restore)
-
-        baseline_logits = _classify_from_tokens(model, module_tokens)
-        y_tensor = torch.from_numpy(y_batch).to(device)
-        batch_idx = torch.arange(len(y_batch), device=device)
-        baseline_correct_logits = baseline_logits[
-            batch_idx, y_tensor
-        ].cpu().numpy()
-
-        for pi in range(num_modules):
-            masked_tokens = module_tokens.clone()
-            masked_tokens[:, pi, :] = 0.0
-            masked_logits = _classify_from_tokens(model, masked_tokens)
-            masked_correct_logits = masked_logits[
-                batch_idx, y_tensor
-            ].cpu().numpy()
-
-            delta = baseline_correct_logits - masked_correct_logits
-            for cls_id in range(num_classes):
-                mask_cls = y_batch == cls_id
-                if mask_cls.any():
-                    importance[cls_id, pi] += delta[mask_cls].sum()
-
-        for cls_id in range(num_classes):
-            class_counts[cls_id] += (y_batch == cls_id).sum()
-
-    class_counts[class_counts == 0] = 1.0
-    importance /= class_counts[:, None]
-    return importance
-
-
-@torch.no_grad()
-def compute_baseline_accuracy(
-    model: nn.Module,
-    dataloader: DataLoader,
-    preprocessor,
-    reindex_state: dict,
-    device: torch.device,
-) -> Tuple[float, np.ndarray, np.ndarray]:
-    """Compute baseline accuracy on the test set."""
-    preds_all, labels_all = [], []
-    for x_batch, y_batch in dataloader:
-        x_batch = x_batch.to(device)
-        x_model = preprocessor.reindex_tensor_to_model_space(
-            x_batch, reindex_state=reindex_state
-        )
-        logits = model(x_model)
-        preds_all.append(logits.argmax(dim=-1).cpu().numpy())
-        labels_all.append(y_batch.numpy())
-    preds = np.concatenate(preds_all)
-    labels = np.concatenate(labels_all)
-    acc = (preds == labels).mean()
-    return float(acc), preds, labels
-
-
-def _classify_from_tokens(model: nn.Module, module_tokens: torch.Tensor) -> torch.Tensor:
-    """Run pooler + classifier on module tokens."""
-    if model.head_type == "module_mean":
-        cell_emb = module_tokens.mean(dim=1)
-    elif model.head_type == "max":
-        cell_emb = module_tokens.amax(dim=1)
-    elif model.head_type == "attn_pool":
-        cell_emb, _ = model.pooler(module_tokens)
-    elif model.head_type == "gene_mean":
-        gene_tokens = model._biomod_to_gene(module_tokens)
-        cell_emb = gene_tokens.mean(dim=1)
-    else:
-        raise ValueError(f"Unsupported head_type: {model.head_type}")
-    cell_emb = model.output_dropout(cell_emb)
-    return model.classifier(cell_emb)
-
-
-@torch.no_grad()
-def compute_importance_pre_encoder(
-    model: nn.Module,
-    dataloader: DataLoader,
-    preprocessor,
-    reindex_state: dict,
-    num_modules: int,
-    num_classes: int,
-    device: torch.device,
-    max_cells: int = 0,
-    return_raw: bool = False,
-) -> np.ndarray:
-    """Accurate mode: mask module input before encoder (slow).
-
-    max_cells: if > 0, stop after processing this many cells (for subsampling).
-               if 0, process all cells in the dataloader.
-    return_raw: if True, return (importance_sum, class_counts) without normalizing.
-                Used for DDP aggregation.
-    """
-    importance = np.zeros((num_classes, num_modules), dtype=np.float64)
-    class_counts = np.zeros(num_classes, dtype=np.float64)
-    total_cells = 0
-
-    for x_batch, y_batch in dataloader:
-        if max_cells > 0 and total_cells >= max_cells:
-            break
-        x_batch = x_batch.to(device)
-        y_batch = y_batch.numpy()
-        x_model = preprocessor.reindex_tensor_to_model_space(
-            x_batch, reindex_state=reindex_state
-        )
-
-        baseline_logits = model(x_model)
-        y_tensor = torch.from_numpy(y_batch).to(device)
-        batch_idx = torch.arange(len(y_batch), device=device)
-        baseline_correct_logits = baseline_logits[
-            batch_idx, y_tensor
-        ].cpu().numpy()
-
-        module_sets = model.module_sets
-        for pi in range(num_modules):
-            x_masked = x_model.clone()
-            gene_indices = module_sets[pi]
-            x_masked[:, gene_indices] = 0.0
-            masked_logits = model(x_masked)
-            masked_correct_logits = masked_logits[
-                batch_idx, y_tensor
-            ].cpu().numpy()
-
-            delta = baseline_correct_logits - masked_correct_logits
-            for cls_id in range(num_classes):
-                mask_cls = y_batch == cls_id
-                if mask_cls.any():
-                    importance[cls_id, pi] += delta[mask_cls].sum()
-
-        for cls_id in range(num_classes):
-            class_counts[cls_id] += (y_batch == cls_id).sum()
-        total_cells += len(y_batch)
-
-    if return_raw:
-        return importance, class_counts
-
-    class_counts[class_counts == 0] = 1.0
-    importance /= class_counts[:, None]
-    return importance
-
-
-# ─── Enrichment Analysis ────────────────────────────────────────────────────
-
-def download_panglaodb_markers(cache_dir: Path) -> pd.DataFrame:
-    """Download PanglaoDB marker gene database."""
-    cache_file = cache_dir / "PanglaoDB_markers.tsv"
-    if cache_file.exists():
-        logger.info("Loading cached PanglaoDB markers from %s", cache_file)
-        return pd.read_csv(cache_file, sep="\t")
-
-    url = "https://panglaodb.se/markers/PanglaoDB_markers_27_Mar_2020.tsv.gz"
-    logger.info("Downloading PanglaoDB markers from %s", url)
-    import urllib.request
-    import gzip
-
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    gz_path = cache_dir / "PanglaoDB_markers.tsv.gz"
-    urllib.request.urlretrieve(url, gz_path)
-    with gzip.open(gz_path, "rt") as f_in:
-        content = f_in.read()
-    with open(cache_file, "w") as f_out:
-        f_out.write(content)
-    gz_path.unlink()
-
-    df = pd.read_csv(cache_file, sep="\t")
-    logger.info("Downloaded PanglaoDB: %d entries", len(df))
-    return df
-
-
-def build_marker_dict(
-    panglaodb_df: pd.DataFrame,
-    species: str = "Hs",
-) -> Dict[str, set]:
-    """Build {cell_type: set(gene_symbols)} from PanglaoDB."""
-    if "species" in panglaodb_df.columns:
-        df = panglaodb_df[panglaodb_df["species"].str.contains(species, na=False)]
-    else:
-        df = panglaodb_df
-
-    col_cell = "cell type" if "cell type" in df.columns else "celltype"
-    col_gene = (
-        "official gene symbol" if "official gene symbol" in df.columns else "gene"
-    )
-
-    marker_dict = {}
-    for ct, group in df.groupby(col_cell):
-        genes = set(group[col_gene].dropna().str.strip().str.upper())
-        if genes:
-            marker_dict[ct.strip()] = genes
-    logger.info("Built marker dict: %d cell types", len(marker_dict))
-    return marker_dict
-
-
-def hypergeometric_enrichment(
-    query_genes: set,
-    marker_genes: set,
-    background_size: int,
-) -> Tuple[float, int, int]:
-    """One-sided hypergeometric test for over-representation."""
-    M = background_size
-    n = len(marker_genes)
-    N = len(query_genes)
-    k = len(query_genes & marker_genes)
-    if N == 0 or n == 0:
-        return 1.0, k, N
-    pval = stats.hypergeom.sf(k - 1, M, n, N)
-    return float(pval), k, N
-
-
-def run_enrichment(
-    importance: np.ndarray,
-    module_names: List[str],
-    module_gene_symbols: Dict[str, List[str]],
-    label_classes: List[str],
-    marker_dict: Dict[str, set],
-    top_k: int = 20,
-    background_size: int = 27874,
-) -> pd.DataFrame:
-    """Run hypergeometric enrichment for top-k modules per cell type."""
-    results = []
-    num_classes = importance.shape[0]
-
-    for cls_id in range(num_classes):
-        cls_name = label_classes[cls_id]
-        ranked_indices = np.argsort(-importance[cls_id])
-        top_indices = ranked_indices[:top_k]
-
-        query_genes = set()
-        top_module_names = []
-        for pi in top_indices:
-            pname = module_names[pi]
-            top_module_names.append(pname)
-            genes = module_gene_symbols.get(pname, [])
-            query_genes.update(g.upper() for g in genes)
-
-        for db_ct, db_markers in marker_dict.items():
-            pval, overlap_count, query_size = hypergeometric_enrichment(
-                query_genes, db_markers, background_size
-            )
-            results.append({
-                "model_cell_type": cls_name,
-                "db_cell_type": db_ct,
-                "top_k": top_k,
-                "query_genes": query_size,
-                "db_markers": len(db_markers),
-                "overlap": overlap_count,
-                "pvalue": pval,
-                "top_modules": ";".join(top_module_names),
-            })
-
-    df = pd.DataFrame(results)
-    df["pvalue_adj"] = np.minimum(
-        df["pvalue"] * len(marker_dict), 1.0
-    )
-    return df
-
-
-# ─── Module Gene Mapping ─────────────────────────────────────────────────────
-
-def build_biomod_gene_mapping(
-    module_data_dir: str,
-    gene_list_path: str,
-) -> Tuple[List[str], Dict[str, List[str]]]:
-    """Return (module_names, {module_name: [gene_symbols]})."""
-    merged_path = Path(module_data_dir) / "merged_biomods.json"
-    with open(merged_path, "r") as f:
-        merged = json.load(f)
-
-    module_names = list(merged.keys())
-    module_gene_symbols = {}
-    for pname, genes in merged.items():
-        if genes and isinstance(genes[0], str):
-            module_gene_symbols[pname] = genes
-        else:
-            with open(gene_list_path, "r") as f:
-                gene_list = [line.strip() for line in f if line.strip()]
-            symbols = [gene_list[idx] for idx in genes if idx < len(gene_list)]
-            module_gene_symbols[pname] = symbols
-
-    return module_names, module_gene_symbols
-
-
-# ─── Visualization ──────────────────────────────────────────────────────────
-
-def plot_importance_heatmap(
-    importance: np.ndarray,
-    label_classes: List[str],
-    output_path: Path,
-):
-    """Plot cell-type × module importance heatmap."""
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-
-    fig, ax = plt.subplots(figsize=(16, 6))
-    im = ax.imshow(importance, aspect="auto", cmap="RdBu_r")
-    ax.set_xlabel("Module index")
-    ax.set_ylabel("Cell type")
-    ax.set_yticks(range(len(label_classes)))
-    ax.set_yticklabels(label_classes, fontsize=8)
-    plt.colorbar(im, ax=ax, label="ΔLogit (importance)")
-    plt.title("Module Importance per Cell Type")
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    plt.close()
-    logger.info("Saved heatmap: %s", output_path)
-
-
-def plot_top_biomods_bar(
-    importance: np.ndarray,
-    label_classes: List[str],
-    module_names: List[str],
-    output_dir: Path,
-    top_k: int = 15,
-):
-    """Bar plot of top-k modules for each cell type."""
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-
-    n_classes = len(label_classes)
-    cols = min(3, n_classes)
-    rows = (n_classes + cols - 1) // cols
-    fig, axes = plt.subplots(rows, cols, figsize=(6 * cols, 4 * rows))
-    axes = np.array(axes).flatten()
-
-    for cls_id in range(n_classes):
-        ax = axes[cls_id]
-        ranked = np.argsort(-importance[cls_id])[:top_k]
-        values = importance[cls_id, ranked]
-        names = [module_names[i][:25] for i in ranked]
-        ax.barh(range(top_k), values[::-1])
-        ax.set_yticks(range(top_k))
-        ax.set_yticklabels(names[::-1], fontsize=6)
-        ax.set_title(label_classes[cls_id], fontsize=9)
-        ax.set_xlabel("ΔLogit")
-
-    for i in range(n_classes, len(axes)):
-        axes[i].set_visible(False)
-
-    plt.tight_layout()
-    plt.savefig(output_dir / "top_modules_bar.png", dpi=150, bbox_inches="tight")
-    plt.close()
-    logger.info("Saved bar plots: %s", output_dir / "top_modules_bar.png")
-
-
-# ─── Argument Parser ────────────────────────────────────────────────────────
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Module Importance + Enrichment")
-    p.add_argument("--finetune_checkpoint", required=True,
-                   help="Path to best_model.pt from cell annotation finetuning")
-    p.add_argument("--config_pkl", required=True,
-                   help="Path to pretrained model .pkl config")
-    p.add_argument("--biomod_data_dir", "--module_data_dir", required=True,
-                   help="Directory containing merged_biomods.json", dest='module_data_dir')
-    p.add_argument("--gene_list_path", required=True,
-                   help="Path to selected_genes_27k.txt")
-    p.add_argument("--gene_mapping_path", required=True)
-    p.add_argument("--gene2vec_data_path", required=True)
-    p.add_argument("--gene2ens_file", required=True)
-    p.add_argument("--data_path", required=True,
-                   help="Path to h5ad dataset (e.g. Zheng68K.h5ad)")
-    p.add_argument("--label_column", default="celltype")
-    p.add_argument("--gene_name_column", default="auto")
-    p.add_argument("--head_type", default="attn_pool",
-                   help="Fallback head_type if not in checkpoint")
-    p.add_argument("--mask_mode", default="post_encoder",
-                   choices=["pre_encoder", "post_encoder"])
-    p.add_argument("--importance_cells", type=int, default=0,
-                   help="Max cells for pre_encoder mode. 0 = use all cells.")
-    p.add_argument("--top_k", type=int, nargs="+", default=[10, 20],
-                   help="Top-k modules for enrichment")
-    p.add_argument("--batch_size", type=int, default=64)
-    p.add_argument("--test_ratio", type=float, default=0.2)
-    p.add_argument("--seed", type=int, default=2021)
-    p.add_argument("--marker_db_dir", default=None,
-                   help="Directory to cache marker DB files")
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--device", default="cuda")
-    p.add_argument("--num_workers", type=int, default=0,
-                   help="DataLoader num_workers")
-    return p.parse_args()
-
-
-# ─── Main ───────────────────────────────────────────────────────────────────
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    rank, world_size, device, use_ddp = setup_ddp()
-    logger.info("Device: %s (rank %d/%d, ddp=%s)", device, rank, world_size, use_ddp)
-
-    for attr in ("finetune_checkpoint", "config_pkl", "module_data_dir",
-                 "gene_list_path", "gene_mapping_path", "gene2vec_data_path",
-                 "gene2ens_file", "data_path"):
-        val = getattr(args, attr, None)
-        if val is not None and not Path(val).is_absolute():
-            setattr(args, attr, str(Path(val).resolve()))
-
-    logger.info("Loading finetuned model...")
-    model, biomod_data, label_classes, head_type = load_finetuned_model(args, device)
-    num_modules = len(biomod_data)
-    num_classes = len(label_classes)
-
-    logger.info("Loading test data...")
-    X_test, y_test, le, preprocessor, reindex_state = load_test_data(
-        args, label_classes, device
-    )
-
-    dataset = CellDataset(X_test, y_test)
-
-    # DDP: split cells across ranks for pre_encoder mode
-    if use_ddp and args.mask_mode == "pre_encoder":
-        total = len(dataset)
-        per_rank = total // world_size
-        start = rank * per_rank
-        end = start + per_rank if rank < world_size - 1 else total
-        indices = list(range(start, end))
-        rank_dataset = Subset(dataset, indices)
-        logger.info("Rank %d: cells [%d, %d) (%d cells)",
-                    rank, start, end, len(indices))
-    else:
-        rank_dataset = dataset
-
-    dataloader = DataLoader(
-        rank_dataset, batch_size=args.batch_size, shuffle=False,
-        num_workers=args.num_workers,
-    )
-
-    if rank == 0:
-        full_dataloader = DataLoader(
-            dataset, batch_size=args.batch_size, shuffle=False,
-            num_workers=args.num_workers,
-        )
-        logger.info("Computing baseline accuracy...")
-        baseline_acc, _, _ = compute_baseline_accuracy(
-            model, full_dataloader, preprocessor, reindex_state, device
-        )
-        logger.info("Baseline test accuracy: %.4f", baseline_acc)
-    else:
-        baseline_acc = 0.0
-
-    logger.info("Computing module importance (mode=%s)...", args.mask_mode)
-    if args.mask_mode == "post_encoder":
-        importance = compute_importance_post_encoder(
-            model, dataloader, preprocessor, reindex_state,
-            num_modules, num_classes, device,
-        )
-    else:
-        max_cells = args.importance_cells if args.importance_cells > 0 else 0
-        if use_ddp:
-            imp_raw, cls_counts = compute_importance_pre_encoder(
-                model, dataloader, preprocessor, reindex_state,
-                num_modules, num_classes, device,
-                max_cells=max_cells,
-                return_raw=True,
-            )
-            imp_tensor = torch.from_numpy(imp_raw).to(device)
-            cnt_tensor = torch.from_numpy(cls_counts).to(device)
-            dist.all_reduce(imp_tensor, op=dist.ReduceOp.SUM)
-            dist.all_reduce(cnt_tensor, op=dist.ReduceOp.SUM)
-            cls_counts_total = cnt_tensor.cpu().numpy()
-            cls_counts_total[cls_counts_total == 0] = 1.0
-            importance = imp_tensor.cpu().numpy() / cls_counts_total[:, None]
-            dist.barrier()
-        else:
-            importance = compute_importance_pre_encoder(
-                model, dataloader, preprocessor, reindex_state,
-                num_modules, num_classes, device,
-                max_cells=max_cells,
-            )
-
-    # Only rank 0 does enrichment, plotting, saving
-    if rank == 0:
-        logger.info("Building module-gene mapping...")
-        module_names, module_gene_symbols = build_biomod_gene_mapping(
-            args.module_data_dir, args.gene_list_path
-        )
-
-        np.save(output_dir / "importance_matrix.npy", importance)
-        imp_df = pd.DataFrame(
-            importance, index=label_classes, columns=module_names
-        )
-        imp_df.to_csv(output_dir / "importance_matrix.csv")
-
-        top_modules_info = {}
-        for cls_id in range(num_classes):
-            ranked = np.argsort(-importance[cls_id])[:max(args.top_k)]
-            top_modules_info[label_classes[cls_id]] = [
-                {
-                    "rank": r + 1,
-                    "module": module_names[pi],
-                    "importance": float(importance[cls_id, pi]),
-                    "genes": module_gene_symbols.get(
-                        module_names[pi], []
-                    )[:10],
-                }
-                for r, pi in enumerate(ranked)
-            ]
-        with open(output_dir / "top_modules_per_celltype.json", "w") as f:
-            json.dump(top_modules_info, f, indent=2)
-
-        logger.info("Running enrichment analysis...")
-        marker_cache = (
-            Path(args.marker_db_dir)
-            if args.marker_db_dir
-            else output_dir / "marker_cache"
-        )
-        panglaodb_df = download_panglaodb_markers(marker_cache)
-        marker_dict = build_marker_dict(panglaodb_df)
-
-        all_enrichment = []
-        for k in args.top_k:
-            enrich_df = run_enrichment(
-                importance, module_names, module_gene_symbols,
-                label_classes, marker_dict, top_k=k,
-            )
-            all_enrichment.append(enrich_df)
-        enrichment_df = pd.concat(all_enrichment, ignore_index=True)
-        enrichment_df.to_csv(
-            output_dir / "enrichment_results.csv", index=False
-        )
-
-        sig_hits = enrichment_df[enrichment_df["pvalue_adj"] < 0.05]
-        logger.info(
-            "Enrichment: %d significant hits (adj p < 0.05) out of %d tests",
-            len(sig_hits), len(enrichment_df),
-        )
-
-        logger.info("Generating plots...")
-        plot_importance_heatmap(
-            importance, label_classes, output_dir / "heatmap_importance.png"
-        )
-        plot_top_biomods_bar(
-            importance, label_classes, module_names, output_dir
-        )
-
-        metrics = {
-            "baseline_accuracy": baseline_acc,
-            "mask_mode": args.mask_mode,
-            "num_modules": num_modules,
-            "num_classes": num_classes,
-            "num_test_cells": len(y_test),
-            "top_k_values": args.top_k,
-            "enrichment_significant_hits": int(len(sig_hits)),
-            "enrichment_total_tests": int(len(enrichment_df)),
-        }
-
-        for cls_id in range(num_classes):
-            cls_name = label_classes[cls_id]
-            cls_sig = sig_hits[
-                sig_hits["model_cell_type"] == cls_name
-            ]
-            if len(cls_sig) > 0:
-                best = cls_sig.sort_values("pvalue_adj").iloc[0]
-                metrics[f"best_match_{cls_name}"] = {
-                    "db_cell_type": best["db_cell_type"],
-                    "pvalue_adj": float(best["pvalue_adj"]),
-                    "overlap": int(best["overlap"]),
-                }
-
-        with open(output_dir / "metrics.json", "w") as f:
-            json.dump(metrics, f, indent=2)
-
-        logger.info("Done. Results saved to %s", output_dir)
-
-    cleanup_ddp(use_ddp)
-
-
-if __name__ == "__main__":
-    main()
-

BioMod/interpretability/biomod_importance/biomod_importance_publication.py

@@ -1,382 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Module Importance — Publication-Quality Figure Generation
-
-Generates all figures from saved experiment results (CPU only).
-Run after masking_ablation.py and random_baseline.py have produced results.
-
-Figures:
-  1. Ablation curves (overall + per-class)
-  2. Permutation null distributions
-  3. Marker gene rank dot plot (from marker_gene_verification.py)
-  4. Improved importance heatmap with clustering + annotations
-  5. Enrichment summary heatmap
-"""
-
-import argparse
-import json
-import sys
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-import seaborn as sns
-from scipy.cluster.hierarchy import dendrogram, linkage
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-sys.path.insert(0, str(SCRIPT_DIR))
-
-from pub_utils import (
-    CELL_TYPE_COLORS,
-    MODULE_FAMILY_COLORS,
-    get_biomod_families,
-    load_importance_results,
-    load_biomod_gene_mapping,
-    savefig,
-    setup_publication_style,
-)
-
-
-def plot_ablation_curves_overall(ablation_dir, output_dir):
-    """Fig 1: Overall ablation curves — top-masked only."""
-    data = np.load(ablation_dir / "ablation_curves.npz")
-    k_values = data["k_values"]
-
-    with open(ablation_dir / "ablation_metrics.json") as f:
-        metrics = json.load(f)
-    baseline = metrics["baseline_accuracy"]
-
-    setup_publication_style()
-    fig, ax = plt.subplots(figsize=(6, 4))
-
-    ax.axhline(baseline, color="gray", linestyle=":", linewidth=1.0,
-               label=f"Baseline ({baseline:.1%})")
-
-    ax.plot(k_values, data["overall_top"], "o-", color="#E41A1C",
-            markersize=6, linewidth=2, label="Top modules masked")
-
-    for i, (k, acc) in enumerate(zip(k_values, data["overall_top"])):
-        drop = baseline - acc
-        if drop > 0.02:
-            ax.annotate(f"−{drop:.0%}", (k, acc),
-                        textcoords="offset points", xytext=(0, -14),
-                        fontsize=7, ha="center", color="#E41A1C")
-
-    ax.set_xlabel("Number of modules masked")
-    ax.set_ylabel("Classification accuracy")
-    ax.set_title("Module Masking Ablation (Top Importance Modules)")
-    ax.legend(loc="upper right", framealpha=0.9)
-    ax.set_xscale("log")
-    ax.set_xticks(k_values)
-    ax.set_xticklabels([str(k) for k in k_values])
-    ax.set_ylim(0.4, 0.9)
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig1_ablation_curves")
-    return True
-
-
-def plot_ablation_curves_per_class(ablation_dir, output_dir):
-    """Fig S1: Per-cell-type ablation curves."""
-    data = np.load(ablation_dir / "ablation_curves.npz")
-    k_values = data["k_values"]
-
-    with open(ablation_dir / "ablation_metrics.json") as f:
-        metrics = json.load(f)
-    label_classes = metrics["label_classes"]
-    num_classes = len(label_classes)
-
-    setup_publication_style()
-    ncols = 4
-    nrows = (num_classes + ncols - 1) // ncols
-    fig, axes = plt.subplots(nrows, ncols, figsize=(16, 3.5 * nrows))
-    axes = axes.flatten()
-
-    for cls_idx in range(num_classes):
-        ax = axes[cls_idx]
-        color = CELL_TYPE_COLORS.get(label_classes[cls_idx], "#333333")
-
-        ax.plot(k_values, data["per_class_top"][cls_idx], "o-",
-                color="#E41A1C", markersize=3, label="Top")
-        ax.plot(k_values, data["per_class_bottom"][cls_idx], "s-",
-                color="#377EB8", markersize=3, label="Bottom")
-        ax.plot(k_values, data["per_class_random_mean"][cls_idx], "^-",
-                color="#4DAF4A", markersize=3, label="Random")
-        ax.fill_between(
-            k_values,
-            data["per_class_random_mean"][cls_idx] - data["per_class_random_std"][cls_idx],
-            data["per_class_random_mean"][cls_idx] + data["per_class_random_std"][cls_idx],
-            alpha=0.15, color="#4DAF4A",
-        )
-
-        ax.set_title(label_classes[cls_idx], fontsize=9, color=color)
-        ax.set_xscale("log")
-        ax.set_ylim(0, 1.05)
-        if cls_idx == 0:
-            ax.legend(fontsize=7, loc="lower left")
-
-    for idx in range(num_classes, len(axes)):
-        axes[idx].set_visible(False)
-
-    fig.supxlabel("Number of modules masked", fontsize=10)
-    fig.supylabel("Accuracy", fontsize=10)
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig_s1_ablation_per_class")
-    return True
-
-
-def plot_permutation_null(permutation_dir, output_dir):
-    """Fig 2: Permutation null distributions."""
-    setup_publication_style()
-
-    metrics_path = permutation_dir / "permutation_metrics.json"
-    if not metrics_path.exists():
-        return False
-
-    with open(metrics_path) as f:
-        metrics = json.load(f)
-
-    has_modules = (
-        (permutation_dir / "null_module_shuffle.npy").exists()
-        and "shuffle_modules" in metrics
-    )
-    has_labels = (
-        (permutation_dir / "null_label_shuffle.npy").exists()
-        and "shuffle_labels" in metrics
-    )
-
-    if not has_modules and not has_labels:
-        return False
-
-    n_panels = has_modules + has_labels
-    fig, axes = plt.subplots(1, n_panels, figsize=(5 * n_panels, 4))
-    if n_panels == 1:
-        axes = [axes]
-
-    panel_idx = 0
-
-    if has_modules:
-        null_data = np.load(permutation_dir / "null_module_shuffle.npy")
-        null_max = null_data.max(axis=(1, 2))
-        stats = metrics["shuffle_modules"]
-        real_max = stats["real_global_max"]
-
-        ax = axes[panel_idx]
-        ax.hist(null_max, bins=20, color="#377EB8", alpha=0.7, edgecolor="white")
-        ax.axvline(real_max, color="#E41A1C", linewidth=2, linestyle="--",
-                   label=f"Observed ({real_max:.3f})")
-        ax.set_xlabel("Max importance (global)")
-        ax.set_ylabel("Count")
-        ax.set_title(f"Module-Shuffle Null (n={stats['n_permutations']})\n"
-                     f"p = {stats['pvalue_global_max']:.4f}")
-        ax.legend()
-        panel_idx += 1
-
-    if has_labels:
-        null_data = np.load(permutation_dir / "null_label_shuffle.npy")
-        null_max = null_data.max(axis=(1, 2))
-        stats = metrics["shuffle_labels"]
-        real_max = stats["real_global_max"]
-
-        ax = axes[panel_idx]
-        ax.hist(null_max, bins=30, color="#4DAF4A", alpha=0.7, edgecolor="white")
-        ax.axvline(real_max, color="#E41A1C", linewidth=2, linestyle="--",
-                   label=f"Observed ({real_max:.3f})")
-        ax.set_xlabel("Max importance (global)")
-        ax.set_ylabel("Count")
-        ax.set_title(f"Label-Shuffle Null (n={stats['n_permutations']})\n"
-                     f"p = {stats['pvalue_global_max']:.4f}")
-        ax.legend()
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig2_permutation_null")
-    return True
-
-
-def plot_importance_heatmap_clustered(results_dir, module_data_dir, output_dir):
-    """Fig 4: Improved importance heatmap with clustering and annotations."""
-    setup_publication_style()
-
-    data = load_importance_results(results_dir)
-    importance = data["importance"]
-    label_classes = data["label_classes"]
-    module_names = data["module_names"]
-
-    families = get_biomod_families(module_names)
-    family_colors = [MODULE_FAMILY_COLORS.get(f, "#CCCCCC") for f in families]
-
-    # Cluster modules
-    linkage_modules = linkage(importance.T, method="ward")
-    from scipy.cluster.hierarchy import leaves_list
-    module_order = leaves_list(linkage_modules)
-
-    importance_ordered = importance[:, module_order]
-    families_ordered = [families[i] for i in module_order]
-    family_colors_ordered = [family_colors[i] for i in module_order]
-
-    from matplotlib.patches import Patch as LegendHandle
-    fig = plt.figure(figsize=(14, 6))
-    gs = fig.add_gridspec(
-        2, 2, height_ratios=[1, 12], width_ratios=[50, 1],
-        hspace=0.05, wspace=0.02,
-    )
-
-    # Top bar: module family annotation
-    ax_bar = fig.add_subplot(gs[0, 0])
-    for i, color in enumerate(family_colors_ordered):
-        ax_bar.axvspan(i - 0.5, i + 0.5, color=color, alpha=0.8)
-    ax_bar.set_xlim(-0.5, len(module_order) - 0.5)
-    ax_bar.set_yticks([])
-    ax_bar.set_xticks([])
-    ax_bar.spines[:].set_visible(False)
-
-    # Legend for families — placed above the bar using figure coordinates
-    legend_modules = [
-        LegendHandle(facecolor=c, label=f) for f, c in MODULE_FAMILY_COLORS.items()
-    ]
-    ax_bar.legend(
-        handles=legend_modules, loc="lower center",
-        bbox_to_anchor=(0.5, 1.05), ncol=4, fontsize=8,
-        frameon=False,
-    )
-
-    # Main heatmap
-    ax_heat = fig.add_subplot(gs[1, 0])
-    vmax = np.percentile(np.abs(importance), 98)
-    im = ax_heat.imshow(
-        importance_ordered, aspect="auto", cmap="RdBu_r",
-        vmin=-vmax, vmax=vmax, interpolation="nearest",
-    )
-    ax_heat.set_yticks(range(len(label_classes)))
-    ax_heat.set_yticklabels(label_classes, fontsize=9)
-    ax_heat.set_xlabel("Modules (hierarchically clustered)")
-    ax_heat.set_xticks([])
-
-    # Colorbar in its own subplot area
-    ax_cbar = fig.add_subplot(gs[1, 1])
-    cbar = fig.colorbar(im, cax=ax_cbar)
-    cbar.set_label("ΔLogit (importance)", fontsize=9)
-
-    # Hide unused top-right cell
-    ax_empty = fig.add_subplot(gs[0, 1])
-    ax_empty.set_visible(False)
-
-    fig.suptitle("Module Importance Matrix (hierarchically clustered)", fontsize=11, y=0.98)
-    savefig(fig, output_dir / "fig4_importance_heatmap_clustered")
-    return True
-
-
-def plot_enrichment_summary(results_dir, output_dir):
-    """Fig 5: Enrichment summary heatmap."""
-    setup_publication_style()
-
-    data = load_importance_results(results_dir)
-    if "enrichment" not in data:
-        return False
-
-    enrichment = data["enrichment"]
-    label_classes = data["label_classes"]
-
-    sig = enrichment[enrichment["pvalue_adj"] < 0.05].copy()
-    if sig.empty:
-        return False
-
-    # Get top db_cell_types per model cell type
-    top_matches = []
-    for ct in label_classes:
-        ct_sig = sig[sig["model_cell_type"] == ct].sort_values("pvalue_adj")
-        top_matches.append(ct_sig.head(5))
-    top_df = pd.concat(top_matches)
-
-    db_types = top_df["db_cell_type"].unique()
-
-    # Build matrix
-    matrix = np.zeros((len(label_classes), len(db_types)))
-    for i, ct in enumerate(label_classes):
-        for j, db in enumerate(db_types):
-            row = enrichment[
-                (enrichment["model_cell_type"] == ct) &
-                (enrichment["db_cell_type"] == db)
-            ]
-            if not row.empty:
-                pval = row["pvalue_adj"].min()
-                matrix[i, j] = -np.log10(max(pval, 1e-50))
-
-    fig, ax = plt.subplots(figsize=(max(10, len(db_types) * 0.5), 5))
-    im = ax.imshow(matrix, aspect="auto", cmap="YlOrRd", interpolation="nearest")
-    ax.set_yticks(range(len(label_classes)))
-    ax.set_yticklabels(label_classes, fontsize=8)
-    ax.set_xticks(range(len(db_types)))
-    ax.set_xticklabels(db_types, fontsize=7, rotation=45, ha="right")
-    ax.set_xlabel("PanglaoDB Cell Types")
-    ax.set_title("Enrichment Significance: -log10(adjusted p-value)")
-
-    cbar = fig.colorbar(im, ax=ax, fraction=0.02, pad=0.01)
-    cbar.set_label("-log10(p_adj)", fontsize=9)
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig5_enrichment_summary")
-    return True
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Generate publication figures")
-    p.add_argument("--results_dir", required=True,
-                   help="Base importance results directory")
-    p.add_argument("--ablation_dir",
-                   help="Ablation results directory")
-    p.add_argument("--permutation_dir",
-                   help="Permutation results directory")
-    p.add_argument("--biomod_data_dir", "--module_data_dir", required=True,
-                   help="Module data directory (merged_biomods.json)", dest='module_data_dir')
-    p.add_argument("--output_dir", required=True)
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    results_dir = Path(args.results_dir)
-    module_data_dir = Path(args.module_data_dir)
-
-    generated = []
-
-    # Fig 4: Improved heatmap (always available)
-    print("Generating Fig 4: Importance heatmap (clustered)...")
-    if plot_importance_heatmap_clustered(results_dir, module_data_dir, output_dir):
-        generated.append("fig4_importance_heatmap_clustered")
-
-    # Fig 5: Enrichment summary (always available)
-    print("Generating Fig 5: Enrichment summary...")
-    if plot_enrichment_summary(results_dir, output_dir):
-        generated.append("fig5_enrichment_summary")
-
-    # Fig 1: Ablation curves (if available)
-    if args.ablation_dir:
-        ablation_dir = Path(args.ablation_dir)
-        if (ablation_dir / "ablation_curves.npz").exists():
-            print("Generating Fig 1: Ablation curves...")
-            if plot_ablation_curves_overall(ablation_dir, output_dir):
-                generated.append("fig1_ablation_curves")
-            print("Generating Fig S1: Per-class ablation...")
-            if plot_ablation_curves_per_class(ablation_dir, output_dir):
-                generated.append("fig_s1_ablation_per_class")
-
-    # Fig 2: Permutation null (if available)
-    if args.permutation_dir:
-        permutation_dir = Path(args.permutation_dir)
-        if (permutation_dir / "permutation_metrics.json").exists():
-            print("Generating Fig 2: Permutation null...")
-            if plot_permutation_null(permutation_dir, output_dir):
-                generated.append("fig2_permutation_null")
-
-    print(f"\nGenerated {len(generated)} figures: {generated}")
-    print(f"Output: {output_dir}")
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/cross_class_similarity.py

@@ -1,247 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Cross-class module importance similarity analysis.
-
-Generates a multi-panel figure showing:
-  A) Spearman correlation of importance profiles across cell types
-  B) Jaccard similarity of top-K important modules
-  C) Shared vs. class-specific module counts (UpSet-style bar)
-
-Runs entirely on saved results (CPU only).
-"""
-
-import argparse
-import sys
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import matplotlib as mpl
-import numpy as np
-import pandas as pd
-from scipy.cluster.hierarchy import linkage, leaves_list
-from scipy.spatial.distance import squareform
-from scipy.stats import spearmanr
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-sys.path.insert(0, str(SCRIPT_DIR))
-
-from pub_utils import (
-    CELL_TYPE_COLORS,
-    load_importance_results,
-    savefig,
-    setup_publication_style,
-)
-
-
-def compute_spearman_corr(importance: np.ndarray, label_classes: list):
-    """Pairwise Spearman correlation between cell-type importance profiles."""
-    n = len(label_classes)
-    corr_mat = np.zeros((n, n))
-    pval_mat = np.ones((n, n))
-    for i in range(n):
-        for j in range(i, n):
-            if np.std(importance[i]) == 0 or np.std(importance[j]) == 0:
-                corr_mat[i, j] = corr_mat[j, i] = np.nan
-                pval_mat[i, j] = pval_mat[j, i] = np.nan
-            else:
-                r, p = spearmanr(importance[i], importance[j])
-                corr_mat[i, j] = corr_mat[j, i] = r
-                pval_mat[i, j] = pval_mat[j, i] = p
-    return corr_mat, pval_mat
-
-
-def compute_jaccard(importance: np.ndarray, K: int):
-    """Jaccard similarity of top-K modules between every pair of cell types."""
-    n = importance.shape[0]
-    top_sets = [set(np.argsort(-importance[i])[:K]) for i in range(n)]
-    jac = np.zeros((n, n))
-    for i in range(n):
-        for j in range(n):
-            union = len(top_sets[i] | top_sets[j])
-            jac[i, j] = len(top_sets[i] & top_sets[j]) / union if union > 0 else 0
-    return jac, top_sets
-
-
-def compute_sharing_stats(importance: np.ndarray, label_classes: list, K: int):
-    """Count how many cell types each top-K module is important for."""
-    n = len(label_classes)
-    top_sets = [set(np.argsort(-importance[i])[:K]) for i in range(n)]
-    all_top = set()
-    for s in top_sets:
-        all_top |= s
-
-    sharing = {}
-    for p in all_top:
-        count = sum(1 for s in top_sets if p in s)
-        sharing[p] = count
-
-    counts = np.array(list(sharing.values()))
-    unique_per_class = {}
-    for i, ct in enumerate(label_classes):
-        unique_per_class[ct] = sum(1 for p in top_sets[i] if sharing[p] == 1)
-
-    return sharing, unique_per_class, counts
-
-
-def plot_cross_class_similarity(results_dir, output_dir, K=50):
-    setup_publication_style()
-
-    data = load_importance_results(results_dir)
-    importance = data["importance"]
-    label_classes = data["label_classes"]
-
-    # --- Compute ---
-    corr_mat, pval_mat = compute_spearman_corr(importance, label_classes)
-    jac_mat, top_sets = compute_jaccard(importance, K)
-    sharing, unique_per_class, sharing_counts = compute_sharing_stats(
-        importance, label_classes, K
-    )
-
-    # --- Cluster order (from Jaccard, excluding NaN rows) ---
-    valid_mask = ~np.isnan(corr_mat).all(axis=1)
-    valid_idx = np.where(valid_mask)[0]
-
-    if len(valid_idx) >= 2:
-        jac_valid = jac_mat[np.ix_(valid_idx, valid_idx)]
-        dist = 1 - jac_valid
-        np.fill_diagonal(dist, 0)
-        link = linkage(squareform(dist), method="average")
-        order_valid = leaves_list(link)
-        display_order = valid_idx[order_valid]
-        # Append NaN rows at the end
-        nan_idx = np.where(~valid_mask)[0]
-        display_order = np.concatenate([display_order, nan_idx])
-    else:
-        display_order = np.arange(len(label_classes))
-
-    # Reorder matrices
-    corr_ordered = corr_mat[np.ix_(display_order, display_order)]
-    jac_ordered = jac_mat[np.ix_(display_order, display_order)]
-    labels_ordered = [label_classes[i] for i in display_order]
-
-    # --- Figure ---
-    fig, axes_all = plt.subplots(1, 3, figsize=(22, 8),
-                                 gridspec_kw={"width_ratios": [1, 1, 0.65], "wspace": 0.45})
-
-    # Panel A: Spearman correlation
-    ax_a = axes_all[0]
-    mask_a = np.zeros_like(corr_ordered, dtype=bool)
-    mask_a[np.triu_indices_from(mask_a, k=1)] = True
-
-    im_a = ax_a.imshow(corr_ordered, cmap="RdBu_r", vmin=-1, vmax=1, aspect="equal")
-    ax_a.set_xticks(range(len(labels_ordered)))
-    ax_a.set_xticklabels(labels_ordered, rotation=45, ha="right", fontsize=7)
-    ax_a.set_yticks(range(len(labels_ordered)))
-    ax_a.set_yticklabels(labels_ordered, fontsize=7)
-
-    # Annotate cells
-    n = len(labels_ordered)
-    for i in range(n):
-        for j in range(n):
-            val = corr_ordered[i, j]
-            if np.isnan(val):
-                ax_a.text(j, i, "—", ha="center", va="center", fontsize=6, color="gray")
-            else:
-                color = "white" if abs(val) > 0.5 else "black"
-                ax_a.text(j, i, f"{val:.2f}", ha="center", va="center",
-                          fontsize=5.5, color=color)
-
-    cb_a = fig.colorbar(im_a, ax=ax_a, fraction=0.046, pad=0.02, shrink=0.85)
-    cb_a.set_label("Spearman ρ", fontsize=9)
-    ax_a.set_title("(A) Importance Profile Correlation", fontsize=10, pad=10)
-
-    # Panel B: Jaccard similarity
-    ax_b = axes_all[1]
-    im_b = ax_b.imshow(jac_ordered, cmap="YlOrRd", vmin=0, vmax=0.35, aspect="equal")
-    ax_b.set_xticks(range(len(labels_ordered)))
-    ax_b.set_xticklabels(labels_ordered, rotation=45, ha="right", fontsize=7)
-    ax_b.set_yticks(range(len(labels_ordered)))
-    ax_b.set_yticklabels(labels_ordered, fontsize=7)
-
-    for i in range(n):
-        for j in range(n):
-            val = jac_ordered[i, j]
-            color = "white" if val > 0.2 else "black"
-            ax_b.text(j, i, f"{val:.2f}", ha="center", va="center",
-                      fontsize=5.5, color=color)
-
-    cb_b = fig.colorbar(im_b, ax=ax_b, fraction=0.046, pad=0.02, shrink=0.85)
-    cb_b.set_label("Jaccard Index", fontsize=9)
-    ax_b.set_title(f"(B) Top-{K} Module Overlap (Jaccard)", fontsize=10, pad=10)
-
-    # Panel C: Sharing distribution bar chart
-    ax_c = axes_all[2]
-
-    # How many of the union of all top-K modules are shared by 1, 2, ..., n classes
-    max_share = sharing_counts.max()
-    bins = np.arange(1, max_share + 2) - 0.5
-    ax_c.hist(sharing_counts, bins=bins, color="#377EB8", edgecolor="white", alpha=0.85)
-    ax_c.set_xlabel("# cell types sharing the module", fontsize=9)
-    ax_c.set_ylabel("# modules", fontsize=9)
-    ax_c.set_title(f"(C) Module Sharing Distribution\n(top-{K} per class)", fontsize=10, pad=10)
-    ax_c.set_xticks(range(1, int(max_share) + 1))
-
-    total_union = len(sharing)
-    unique_total = sum(1 for v in sharing.values() if v == 1)
-    pct = 100 * unique_total / total_union
-    ax_c.annotate(
-        f"{unique_total}/{total_union} modules\n({pct:.0f}%) class-specific",
-        xy=(0.95, 0.92), xycoords="axes fraction",
-        fontsize=8, color="#E41A1C", fontweight="bold",
-        ha="right", va="top",
-    )
-
-    fig.suptitle(
-        "Cross-Class Similarity of Module Importance Profiles",
-        fontsize=12, y=0.98,
-    )
-    fig.subplots_adjust(bottom=0.22, top=0.90, left=0.08, right=0.95)
-    savefig(fig, output_dir / "fig_cross_class_similarity")
-
-    # --- Also save numerical results ---
-    corr_df = pd.DataFrame(corr_mat, index=label_classes, columns=label_classes)
-    corr_df.to_csv(output_dir / "cross_class_spearman_corr.csv")
-
-    jac_df = pd.DataFrame(jac_mat, index=label_classes, columns=label_classes)
-    jac_df.to_csv(output_dir / f"cross_class_jaccard_top{K}.csv")
-
-    unique_df = pd.DataFrame({
-        "cell_type": label_classes,
-        "unique_modules_in_top50": [unique_per_class[ct] for ct in label_classes],
-    })
-    unique_df.to_csv(output_dir / "class_specific_module_counts.csv", index=False)
-
-    print(f"\nSummary:")
-    print(f"  Total unique modules in union of top-{K}: {total_union}")
-    print(f"  Class-specific (appear in only 1 class top-{K}): {unique_total} ({pct:.1f}%)")
-    print(f"  Shared by ≥2 classes: {total_union - unique_total}")
-    print(f"\nSpearman corr range (excl. diagonal & NaN): "
-          f"[{np.nanmin(corr_mat[~np.eye(n, dtype=bool)]):.3f}, "
-          f"{np.nanmax(corr_mat[~np.eye(n, dtype=bool)]):.3f}]")
-    print(f"Jaccard range (excl. diagonal): "
-          f"[{jac_mat[~np.eye(n, dtype=bool)].min():.3f}, "
-          f"{jac_mat[~np.eye(n, dtype=bool)].max():.3f}]")
-
-    return True
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Cross-class module importance similarity")
-    p.add_argument("--results_dir", required=True,
-                   help="Importance results directory (with importance_matrix.npy)")
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--top_k", type=int, default=50,
-                   help="Number of top modules for Jaccard (default: 50)")
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    plot_cross_class_similarity(Path(args.results_dir), output_dir, K=args.top_k)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/family_ablation.py

@@ -1,370 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Experiment ③-D: Family-Level Ablation
-
-Masks entire module families (Hallmarks / Reactome / Hotspot / WGCNA) and measures
-accuracy drop. Two modes:
-  (a) Remove-one-family: mask all modules of a single family
-  (b) Keep-one-family:   mask all modules EXCEPT one family
-
-Requires GPU (model inference). Uses post_encoder masking.
-
-Outputs:
-  - family_ablation_results.json: accuracy for each condition
-  - fig_family_ablation.png/pdf: publication figure
-"""
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-
-import matplotlib as mpl
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-from torch.utils.data import DataLoader
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-INTERP_DIR = SCRIPT_DIR.parent
-BIOMOD_ROOT = INTERP_DIR.parent
-SCLONG_ROOT = BIOMOD_ROOT.parent
-
-for p in (str(SCRIPT_DIR), str(BIOMOD_ROOT), str(SCLONG_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from masking_ablation import (
-    compute_masked_accuracy_post_encoder,
-    encode_all_cells,
-)
-from biomod_importance import (
-    CellDataset,
-    _classify_from_tokens,
-    build_biomod_gene_mapping,
-    load_finetuned_model,
-    load_test_data,
-)
-from pub_utils import (
-    CELL_TYPE_COLORS,
-    savefig,
-    setup_publication_style,
-)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger(__name__)
-
-FAMILY_PREFIXES = {
-    "Hallmarks": "module1_hallmarks",
-    "Reactome": "module2_reactome",
-    "Hotspot": "module3_hotspot",
-    "WGCNA": "module3_processed",
-}
-
-FAMILY_ORDER = ["Hallmarks", "Reactome", "Hotspot", "WGCNA"]
-
-FAMILY_COLORS = {
-    "Hallmarks": "#D95F02",
-    "Reactome": "#1B9E77",
-    "Hotspot": "#7570B3",
-    "WGCNA": "#E7298A",
-}
-
-
-def get_family_indices(module_names):
-    """Return dict: family_name -> list of module indices."""
-    indices = {f: [] for f in FAMILY_ORDER}
-    for i, name in enumerate(module_names):
-        for fam, prefix in FAMILY_PREFIXES.items():
-            if name.startswith(prefix):
-                indices[fam].append(i)
-                break
-    return indices
-
-
-@torch.no_grad()
-def compute_per_class_accuracy(model, all_tokens, all_labels, mask_indices,
-                               label_classes, device, batch_size=256):
-    """Compute overall and per-class accuracy with given modules masked."""
-    n_cells = all_tokens.shape[0]
-    num_classes = len(label_classes)
-    correct_per_class = np.zeros(num_classes)
-    total_per_class = np.zeros(num_classes)
-
-    for start in range(0, n_cells, batch_size):
-        end = min(start + batch_size, n_cells)
-        tokens = all_tokens[start:end].to(device).clone()
-        labels = all_labels[start:end]
-
-        for pi in mask_indices:
-            tokens[:, pi, :] = 0.0
-
-        logits = _classify_from_tokens(model, tokens)
-        preds = logits.argmax(dim=-1).cpu().numpy()
-
-        for cls_idx in range(num_classes):
-            cls_mask = labels == cls_idx
-            correct_per_class[cls_idx] += (preds[cls_mask] == labels[cls_mask]).sum()
-            total_per_class[cls_idx] += cls_mask.sum()
-
-    per_class_acc = np.where(total_per_class > 0,
-                             correct_per_class / total_per_class, 0.0)
-    overall_acc = correct_per_class.sum() / total_per_class.sum()
-    return overall_acc, per_class_acc
-
-
-def run_family_ablation(model, all_tokens, all_labels, module_names,
-                        label_classes, device, batch_size):
-    """Run remove-one and keep-one family ablation experiments."""
-    family_indices = get_family_indices(module_names)
-    all_module_indices = list(range(all_tokens.shape[1]))
-    results = {}
-
-    # Baseline (no masking)
-    logger.info("Computing baseline accuracy...")
-    baseline_overall, baseline_per_class = compute_per_class_accuracy(
-        model, all_tokens, all_labels, [], label_classes, device, batch_size
-    )
-    results["baseline"] = {
-        "overall": float(baseline_overall),
-        "per_class": {ct: float(baseline_per_class[i]) for i, ct in enumerate(label_classes)},
-    }
-    logger.info("Baseline accuracy: %.4f", baseline_overall)
-
-    # Remove-one-family
-    results["remove_one"] = {}
-    for fam in FAMILY_ORDER:
-        mask_indices = family_indices[fam]
-        logger.info("Remove %s (%d modules)...", fam, len(mask_indices))
-        overall, per_class = compute_per_class_accuracy(
-            model, all_tokens, all_labels, mask_indices, label_classes, device, batch_size
-        )
-        results["remove_one"][fam] = {
-            "n_masked": len(mask_indices),
-            "overall": float(overall),
-            "delta_overall": float(baseline_overall - overall),
-            "per_class": {ct: float(per_class[i]) for i, ct in enumerate(label_classes)},
-            "delta_per_class": {ct: float(baseline_per_class[i] - per_class[i])
-                                for i, ct in enumerate(label_classes)},
-        }
-        logger.info("  Accuracy: %.4f (Δ=%.4f)", overall, baseline_overall - overall)
-
-    # Keep-one-family (mask all others)
-    results["keep_one"] = {}
-    for fam in FAMILY_ORDER:
-        keep_indices = set(family_indices[fam])
-        mask_indices = [i for i in all_module_indices if i not in keep_indices]
-        logger.info("Keep only %s (%d modules, mask %d)...",
-                    fam, len(keep_indices), len(mask_indices))
-        overall, per_class = compute_per_class_accuracy(
-            model, all_tokens, all_labels, mask_indices, label_classes, device, batch_size
-        )
-        results["keep_one"][fam] = {
-            "n_kept": len(keep_indices),
-            "n_masked": len(mask_indices),
-            "overall": float(overall),
-            "delta_overall": float(baseline_overall - overall),
-            "per_class": {ct: float(per_class[i]) for i, ct in enumerate(label_classes)},
-            "delta_per_class": {ct: float(baseline_per_class[i] - per_class[i])
-                                for i, ct in enumerate(label_classes)},
-        }
-        logger.info("  Accuracy: %.4f (Δ=%.4f)", overall, baseline_overall - overall)
-
-    return results
-
-
-def plot_family_ablation(results, label_classes, output_dir):
-    """Generate publication figure from ablation results."""
-    setup_publication_style()
-
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
-
-    baseline = results["baseline"]["overall"]
-    x = np.arange(len(FAMILY_ORDER))
-    width = 0.35
-
-    # Panel A: Remove-one-family (overall)
-    remove_accs = [results["remove_one"][f]["overall"] for f in FAMILY_ORDER]
-    remove_deltas = [results["remove_one"][f]["delta_overall"] for f in FAMILY_ORDER]
-    bars1 = ax1.bar(x, remove_accs, width=0.6,
-                    color=[FAMILY_COLORS[f] for f in FAMILY_ORDER],
-                    edgecolor="white", linewidth=0.5)
-    ax1.axhline(baseline, color="gray", linestyle="--", linewidth=1.2,
-                label=f"Baseline ({baseline:.1%})")
-
-    for i, (acc, delta) in enumerate(zip(remove_accs, remove_deltas)):
-        if delta > 0.005:
-            ax1.annotate(f"−{delta:.1%}", (i, acc), textcoords="offset points",
-                         xytext=(0, -14), ha="center", fontsize=8, color="#E41A1C",
-                         fontweight="bold")
-        else:
-            ax1.annotate(f"−{delta:.1%}", (i, acc), textcoords="offset points",
-                         xytext=(0, 6), ha="center", fontsize=8, color="gray")
-
-    ax1.set_xticks(x)
-    ax1.set_xticklabels([f"{f}\n(n={results['remove_one'][f]['n_masked']})"
-                         for f in FAMILY_ORDER], fontsize=9)
-    ax1.set_ylabel("Accuracy", fontsize=10)
-    ax1.set_title("(A) Remove One Family", fontsize=11)
-    ax1.legend(fontsize=8, loc="lower right")
-    ax1.set_ylim(0, min(1.05, baseline * 1.15))
-
-    # Panel B: Keep-one-family (overall)
-    keep_accs = [results["keep_one"][f]["overall"] for f in FAMILY_ORDER]
-    keep_deltas = [results["keep_one"][f]["delta_overall"] for f in FAMILY_ORDER]
-    bars2 = ax2.bar(x, keep_accs, width=0.6,
-                    color=[FAMILY_COLORS[f] for f in FAMILY_ORDER],
-                    edgecolor="white", linewidth=0.5)
-    ax2.axhline(baseline, color="gray", linestyle="--", linewidth=1.2,
-                label=f"Baseline ({baseline:.1%})")
-
-    for i, (acc, delta) in enumerate(zip(keep_accs, keep_deltas)):
-        ax2.annotate(f"{acc:.1%}", (i, acc), textcoords="offset points",
-                     xytext=(0, 6), ha="center", fontsize=8, fontweight="bold")
-
-    ax2.set_xticks(x)
-    ax2.set_xticklabels([f"{f}\n(n={results['keep_one'][f]['n_kept']})"
-                         for f in FAMILY_ORDER], fontsize=9)
-    ax2.set_ylabel("Accuracy", fontsize=10)
-    ax2.set_title("(B) Keep Only One Family", fontsize=11)
-    ax2.legend(fontsize=8, loc="upper right")
-    ax2.set_ylim(0, min(1.05, baseline * 1.15))
-
-    fig.suptitle("Family-Level Ablation: Module Type Contribution to Classification",
-                 fontsize=12, y=1.02)
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig_family_ablation")
-    logger.info("Saved fig_family_ablation")
-
-    # --- Per-class heatmap for remove-one ---
-    fig2, (ax_r, ax_k) = plt.subplots(1, 2, figsize=(14, 6))
-
-    # Remove-one per-class delta heatmap
-    delta_remove = np.zeros((len(label_classes), len(FAMILY_ORDER)))
-    for fi, fam in enumerate(FAMILY_ORDER):
-        for ci, ct in enumerate(label_classes):
-            delta_remove[ci, fi] = results["remove_one"][fam]["delta_per_class"][ct]
-
-    vmax_r = max(np.abs(delta_remove).max(), 0.05)
-    im_r = ax_r.imshow(delta_remove, cmap="RdYlGn_r", aspect="auto",
-                       vmin=-vmax_r, vmax=vmax_r)
-    ax_r.set_xticks(range(len(FAMILY_ORDER)))
-    ax_r.set_xticklabels(FAMILY_ORDER, fontsize=9)
-    ax_r.set_yticks(range(len(label_classes)))
-    ax_r.set_yticklabels(label_classes, fontsize=8)
-    ax_r.set_title("(A) Accuracy Drop: Remove One Family", fontsize=10)
-    for i in range(len(label_classes)):
-        for j in range(len(FAMILY_ORDER)):
-            val = delta_remove[i, j]
-            color = "white" if abs(val) > vmax_r * 0.6 else "black"
-            sign = "+" if val < 0 else ""
-            ax_r.text(j, i, f"{sign}{-val:.1%}" if val != 0 else "0%",
-                      ha="center", va="center", fontsize=7, color=color)
-    cb_r = fig2.colorbar(im_r, ax=ax_r, fraction=0.03, pad=0.02)
-    cb_r.set_label("ΔAccuracy", fontsize=8)
-
-    # Keep-one per-class accuracy heatmap
-    acc_keep = np.zeros((len(label_classes), len(FAMILY_ORDER)))
-    for fi, fam in enumerate(FAMILY_ORDER):
-        for ci, ct in enumerate(label_classes):
-            acc_keep[ci, fi] = results["keep_one"][fam]["per_class"][ct]
-
-    im_k = ax_k.imshow(acc_keep, cmap="YlGn", aspect="auto", vmin=0, vmax=1)
-    ax_k.set_xticks(range(len(FAMILY_ORDER)))
-    ax_k.set_xticklabels(FAMILY_ORDER, fontsize=9)
-    ax_k.set_yticks(range(len(label_classes)))
-    ax_k.set_yticklabels(label_classes, fontsize=8)
-    ax_k.set_title("(B) Per-Class Accuracy: Keep Only One Family", fontsize=10)
-    for i in range(len(label_classes)):
-        for j in range(len(FAMILY_ORDER)):
-            val = acc_keep[i, j]
-            color = "white" if val > 0.6 else "black"
-            ax_k.text(j, i, f"{val:.0%}", ha="center", va="center",
-                      fontsize=7, color=color)
-    cb_k = fig2.colorbar(im_k, ax=ax_k, fraction=0.03, pad=0.02)
-    cb_k.set_label("Accuracy", fontsize=8)
-
-    fig2.suptitle("Family-Level Ablation: Per-Cell-Type Breakdown", fontsize=12, y=1.01)
-    fig2.tight_layout()
-    savefig(fig2, output_dir / "fig_family_ablation_per_class")
-    logger.info("Saved fig_family_ablation_per_class")
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Family-level ablation experiment")
-    p.add_argument("--finetune_checkpoint", required=True)
-    p.add_argument("--config_pkl", required=True)
-    p.add_argument("--biomod_data_dir", "--module_data_dir", required=True, dest='module_data_dir')
-    p.add_argument("--gene_list_path", required=True)
-    p.add_argument("--gene_mapping_path", required=True)
-    p.add_argument("--gene2vec_data_path", required=True)
-    p.add_argument("--gene2ens_file", required=True)
-    p.add_argument("--data_path", required=True)
-    p.add_argument("--label_column", default="celltype")
-    p.add_argument("--head_type", default="attn_pool")
-    p.add_argument("--gene_name_column", default="auto")
-    p.add_argument("--batch_size", type=int, default=256)
-    p.add_argument("--test_ratio", type=float, default=0.2)
-    p.add_argument("--seed", type=int, default=2021)
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--device", default="cuda")
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-    logger.info("Device: %s", device)
-
-    for attr in ("finetune_checkpoint", "config_pkl", "module_data_dir",
-                 "gene_list_path", "gene_mapping_path", "gene2vec_data_path",
-                 "gene2ens_file", "data_path"):
-        val = getattr(args, attr, None)
-        if val is not None and not Path(val).is_absolute():
-            setattr(args, attr, str(Path(val).resolve()))
-
-    logger.info("Loading model...")
-    model, biomod_data, label_classes, head_type = load_finetuned_model(args, device)
-    module_names, _ = build_biomod_gene_mapping(args.module_data_dir, args.gene_list_path)
-    family_indices = get_family_indices(module_names)
-    for fam in FAMILY_ORDER:
-        logger.info("  %s: %d modules", fam, len(family_indices[fam]))
-
-    logger.info("Loading test data...")
-    X_test, y_test, le, preprocessor, reindex_state = load_test_data(
-        args, label_classes, device
-    )
-    dataset = CellDataset(X_test, y_test)
-    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
-
-    logger.info("Encoding all cells...")
-    all_tokens, all_labels = encode_all_cells(
-        model, dataloader, preprocessor, reindex_state, device
-    )
-    logger.info("Encoded %d cells, %d modules", all_tokens.shape[0], all_tokens.shape[1])
-
-    logger.info("Running family ablation...")
-    results = run_family_ablation(
-        model, all_tokens, all_labels, module_names,
-        label_classes, device, args.batch_size
-    )
-
-    # Save raw results
-    with open(output_dir / "family_ablation_results.json", "w") as f:
-        json.dump(results, f, indent=2)
-    logger.info("Saved family_ablation_results.json")
-
-    # Generate figures
-    logger.info("Generating figures...")
-    plot_family_ablation(results, label_classes, output_dir)
-
-    logger.info("Done. Results saved to %s", output_dir)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/family_importance_analysis.py

@@ -1,417 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Family-Level Module Importance Analysis (CPU only)
-
-Analyzes importance at the module-family level (Hallmarks / Reactome / Hotspot / WGCNA):
-  1. Importance distribution per family (violin + box plot)
-  2. Family contribution to top-K important modules (stacked bar)
-  3. Enrichment test: is each family over/under-represented in top-K,  (Fisher's exact)
-
-All analyses run from saved importance_matrix — no GPU needed.
-"""
-
-import argparse
-import json
-import sys
-from pathlib import Path
-
-import matplotlib as mpl
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-from scipy.stats import fisher_exact, kruskal, mannwhitneyu
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-sys.path.insert(0, str(SCRIPT_DIR))
-
-from pub_utils import (
-    CELL_TYPE_COLORS,
-    load_importance_results,
-    savefig,
-    setup_publication_style,
-)
-
-FAMILY_PREFIXES = {
-    "Hallmarks": "module1_hallmarks",
-    "Reactome": "module2_reactome",
-    "Hotspot": "module3_hotspot",
-    "WGCNA": "module3_processed",
-}
-
-FAMILY_ORDER = ["Hallmarks", "Reactome", "Hotspot", "WGCNA"]
-
-FAMILY_COLORS = {
-    "Hallmarks": "#D95F02",
-    "Reactome": "#1B9E77",
-    "Hotspot": "#7570B3",
-    "WGCNA": "#E7298A",
-}
-
-
-def assign_families(module_names):
-    """Assign each module to a family. Returns list of family labels."""
-    families = []
-    for name in module_names:
-        matched = False
-        for fam, prefix in FAMILY_PREFIXES.items():
-            if name.startswith(prefix):
-                families.append(fam)
-                matched = True
-                break
-        if not matched:
-            families.append("Unknown")
-    return families
-
-
-def compute_family_indices(module_names):
-    """Return dict: family_name -> list of column indices."""
-    families = assign_families(module_names)
-    indices = {f: [] for f in FAMILY_ORDER}
-    for i, fam in enumerate(families):
-        if fam in indices:
-            indices[fam].append(i)
-    return indices
-
-
-# ─── Analysis 1: Importance Distribution ────────────────────────────────────
-
-def plot_family_importance_distribution(importance, label_classes, module_names, output_dir):
-    """Violin + box plot of importance scores per family, for each cell type."""
-    setup_publication_style()
-    family_indices = compute_family_indices(module_names)
-
-    ncols = 4
-    nrows = (len(label_classes) + ncols - 1) // ncols
-    fig, axes = plt.subplots(nrows, ncols, figsize=(16, 3.8 * nrows), sharey=False)
-    axes = axes.flatten()
-
-    for cls_idx, ct in enumerate(label_classes):
-        ax = axes[cls_idx]
-        data_list = []
-        positions = []
-        for fi, fam in enumerate(FAMILY_ORDER):
-            vals = importance[cls_idx, family_indices[fam]]
-            data_list.append(vals)
-            positions.append(fi)
-
-        vp = ax.violinplot(data_list, positions=positions, showextrema=False,
-                           showmedians=False, widths=0.7)
-        for fi, body in enumerate(vp["bodies"]):
-            body.set_facecolor(FAMILY_COLORS[FAMILY_ORDER[fi]])
-            body.set_alpha(0.4)
-
-        bp = ax.boxplot(data_list, positions=positions, widths=0.3,
-                        patch_artist=True, showfliers=False)
-        for fi, module in enumerate(bp["boxes"]):
-            module.set_facecolor(FAMILY_COLORS[FAMILY_ORDER[fi]])
-            module.set_alpha(0.7)
-        for element in ["whiskers", "caps", "medians"]:
-            for line in bp[element]:
-                line.set_color("black")
-                line.set_linewidth(0.8)
-
-        color = CELL_TYPE_COLORS.get(ct, "#333")
-        ax.set_title(ct, fontsize=8, fontweight="bold", color=color)
-        ax.set_xticks(positions)
-        ax.set_xticklabels(FAMILY_ORDER, fontsize=7, rotation=30, ha="right")
-        ax.axhline(0, color="gray", linewidth=0.5, linestyle=":")
-
-        if cls_idx % ncols == 0:
-            ax.set_ylabel("ΔLogit", fontsize=8)
-
-    for idx in range(len(label_classes), len(axes)):
-        axes[idx].set_visible(False)
-
-    fig.suptitle("Module Importance Distribution by Family", fontsize=12, y=1.01)
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig1_family_importance_distribution")
-    print("  Saved fig1_family_importance_distribution")
-
-    # --- Kruskal-Wallis test per cell type ---
-    stats_rows = []
-    for cls_idx, ct in enumerate(label_classes):
-        groups = [importance[cls_idx, family_indices[fam]] for fam in FAMILY_ORDER]
-        if np.std(importance[cls_idx]) == 0:
-            stats_rows.append({"cell_type": ct, "H_statistic": np.nan,
-                               "p_value": np.nan, "significant": False})
-            continue
-        H, p = kruskal(*groups)
-        stats_rows.append({
-            "cell_type": ct, "H_statistic": H, "p_value": p,
-            "significant": p < 0.05,
-        })
-
-    stats_df = pd.DataFrame(stats_rows)
-    stats_df.to_csv(output_dir / "kruskal_wallis_family_test.csv", index=False)
-    print("  Saved kruskal_wallis_family_test.csv")
-
-    # --- Pairwise Mann-Whitney U between families (overall mean importance) ---
-    mean_importance = importance.mean(axis=0)
-    pairwise_rows = []
-    for i, fam_a in enumerate(FAMILY_ORDER):
-        for j, fam_b in enumerate(FAMILY_ORDER):
-            if j <= i:
-                continue
-            vals_a = mean_importance[family_indices[fam_a]]
-            vals_b = mean_importance[family_indices[fam_b]]
-            U, p = mannwhitneyu(vals_a, vals_b, alternative="two-sided")
-            pairwise_rows.append({
-                "family_a": fam_a, "family_b": fam_b,
-                "U_statistic": U, "p_value": p,
-                "mean_a": vals_a.mean(), "mean_b": vals_b.mean(),
-            })
-    pairwise_df = pd.DataFrame(pairwise_rows)
-    pairwise_df.to_csv(output_dir / "pairwise_mannwhitney_families.csv", index=False)
-    print("  Saved pairwise_mannwhitney_families.csv")
-
-    return stats_df
-
-
-# ─── Analysis 2: Family Contribution to Top-K ───────────────────────────────
-
-def plot_family_contribution_topk(importance, label_classes, module_names, output_dir,
-                                  k_values=(10, 20, 50)):
-    """Stacked bar chart: fraction of top-K modules from each family."""
-    setup_publication_style()
-    family_indices = compute_family_indices(module_names)
-    families_arr = np.array(assign_families(module_names))
-
-    n_k = len(k_values)
-    fig, axes = plt.subplots(1, n_k, figsize=(5.5 * n_k, 6), sharey=True)
-    if n_k == 1:
-        axes = [axes]
-
-    contribution_data = {}
-
-    for ki, K in enumerate(k_values):
-        ax = axes[ki]
-        counts = {fam: [] for fam in FAMILY_ORDER}
-
-        for cls_idx, ct in enumerate(label_classes):
-            top_k_idx = np.argsort(-importance[cls_idx])[:K]
-            top_families = families_arr[top_k_idx]
-            for fam in FAMILY_ORDER:
-                counts[fam].append(np.sum(top_families == fam))
-
-        contribution_data[K] = counts
-        x = np.arange(len(label_classes))
-        bottom = np.zeros(len(label_classes))
-
-        for fam in FAMILY_ORDER:
-            vals = np.array(counts[fam], dtype=float)
-            pcts = 100.0 * vals / K
-            ax.bar(x, pcts, bottom=bottom, color=FAMILY_COLORS[fam],
-                   label=fam, width=0.7, edgecolor="white", linewidth=0.3)
-            bottom += pcts
-
-        ax.set_xticks(x)
-        ax.set_xticklabels(label_classes, rotation=45, ha="right", fontsize=7)
-        ax.set_ylabel("% of top modules" if ki == 0 else "")
-        ax.set_title(f"Top-{K} Modules", fontsize=10)
-        ax.set_ylim(0, 105)
-
-        if ki == n_k - 1:
-            ax.legend(loc="upper right", fontsize=8, framealpha=0.9)
-
-    fig.suptitle("Family Contribution to Most Important Modules", fontsize=12, y=1.01)
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig2_family_contribution_topk")
-    print("  Saved fig2_family_contribution_topk")
-
-    # Save numerical data
-    rows = []
-    for K, counts in contribution_data.items():
-        for cls_idx, ct in enumerate(label_classes):
-            row = {"cell_type": ct, "top_k": K}
-            for fam in FAMILY_ORDER:
-                row[f"{fam}_count"] = counts[fam][cls_idx]
-                row[f"{fam}_pct"] = 100.0 * counts[fam][cls_idx] / K
-            rows.append(row)
-    contrib_df = pd.DataFrame(rows)
-    contrib_df.to_csv(output_dir / "family_contribution_topk.csv", index=False)
-    print("  Saved family_contribution_topk.csv")
-
-    return contrib_df
-
-
-# ─── Analysis 3: Family Enrichment Test ─────────────────────────────────────
-
-def run_family_enrichment(importance, label_classes, module_names, output_dir, K=50):
-    """Fisher's exact test: is each family over/under-represented in top-K, """
-    setup_publication_style()
-    family_indices = compute_family_indices(module_names)
-    families_arr = np.array(assign_families(module_names))
-    n_total = len(module_names)
-
-    enrichment_rows = []
-    for cls_idx, ct in enumerate(label_classes):
-        top_k_idx = set(np.argsort(-importance[cls_idx])[:K].tolist())
-        for fam in FAMILY_ORDER:
-            fam_set = set(family_indices[fam])
-            n_fam = len(fam_set)
-            in_top_and_fam = len(top_k_idx & fam_set)
-            in_top_not_fam = K - in_top_and_fam
-            not_top_and_fam = n_fam - in_top_and_fam
-            not_top_not_fam = (n_total - K) - not_top_and_fam
-
-            table = [[in_top_and_fam, in_top_not_fam],
-                     [not_top_and_fam, not_top_not_fam]]
-            odds_ratio, p = fisher_exact(table, alternative="two-sided")
-
-            expected = K * n_fam / n_total
-            fold_enrichment = in_top_and_fam / expected if expected > 0 else np.inf
-
-            enrichment_rows.append({
-                "cell_type": ct,
-                "family": fam,
-                "n_family_total": n_fam,
-                "n_in_topK": in_top_and_fam,
-                "expected_in_topK": round(expected, 2),
-                "fold_enrichment": round(fold_enrichment, 3),
-                "odds_ratio": round(odds_ratio, 3),
-                "p_value": p,
-            })
-
-    enrich_df = pd.DataFrame(enrichment_rows)
-
-    # Bonferroni correction
-    n_tests = len(enrich_df)
-    enrich_df["p_adjusted"] = np.minimum(enrich_df["p_value"] * n_tests, 1.0)
-    enrich_df["significant"] = enrich_df["p_adjusted"] < 0.05
-    enrich_df["direction"] = np.where(
-        enrich_df["fold_enrichment"] > 1, "enriched", "depleted"
-    )
-
-    enrich_df.to_csv(output_dir / "family_enrichment_fisher.csv", index=False)
-    print("  Saved family_enrichment_fisher.csv")
-
-    # --- Heatmap: fold enrichment ---
-    pivot = enrich_df.pivot(index="cell_type", columns="family", values="fold_enrichment")
-    pivot = pivot.reindex(index=label_classes, columns=FAMILY_ORDER)
-    sig_pivot = enrich_df.pivot(index="cell_type", columns="family", values="significant")
-    sig_pivot = sig_pivot.reindex(index=label_classes, columns=FAMILY_ORDER)
-
-    fig, ax = plt.subplots(figsize=(7, 6))
-    vmax = max(pivot.max().max(), 2.5)
-    im = ax.imshow(pivot.values, cmap="RdYlBu_r", aspect="auto",
-                   vmin=0, vmax=vmax)
-    ax.set_xticks(range(len(FAMILY_ORDER)))
-    ax.set_xticklabels(FAMILY_ORDER, fontsize=9)
-    ax.set_yticks(range(len(label_classes)))
-    ax.set_yticklabels(label_classes, fontsize=8)
-
-    for i in range(len(label_classes)):
-        for j in range(len(FAMILY_ORDER)):
-            val = pivot.values[i, j]
-            sig = sig_pivot.values[i, j]
-            txt = f"{val:.1f}"
-            if sig:
-                txt += " *"
-            color = "white" if val > vmax * 0.65 else "black"
-            ax.text(j, i, txt, ha="center", va="center", fontsize=8, color=color)
-
-    cb = fig.colorbar(im, ax=ax, fraction=0.03, pad=0.02)
-    cb.set_label("Fold Enrichment", fontsize=9)
-    ax.set_title(f"Family Enrichment in Top-{K} Modules\n(* p_adj < 0.05, Fisher's exact)",
-                 fontsize=10, pad=10)
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig3_family_enrichment_heatmap")
-    print("  Saved fig3_family_enrichment_heatmap")
-
-    return enrich_df
-
-
-# ─── Summary Figure ─────────────────────────────────────────────────────────
-
-def plot_family_summary(importance, label_classes, module_names, output_dir, K=50):
-    """Combined summary: mean importance per family per cell type."""
-    setup_publication_style()
-    family_indices = compute_family_indices(module_names)
-
-    mean_by_family = np.zeros((len(label_classes), len(FAMILY_ORDER)))
-    for cls_idx in range(len(label_classes)):
-        for fi, fam in enumerate(FAMILY_ORDER):
-            mean_by_family[cls_idx, fi] = importance[cls_idx, family_indices[fam]].mean()
-
-    fig, ax = plt.subplots(figsize=(8, 5))
-    x = np.arange(len(label_classes))
-    width = 0.18
-
-    for fi, fam in enumerate(FAMILY_ORDER):
-        offset = (fi - 1.5) * width
-        ax.bar(x + offset, mean_by_family[:, fi], width=width,
-               color=FAMILY_COLORS[fam], label=fam, edgecolor="white", linewidth=0.3)
-
-    ax.set_xticks(x)
-    ax.set_xticklabels(label_classes, rotation=45, ha="right", fontsize=8)
-    ax.set_ylabel("Mean ΔLogit (importance)", fontsize=9)
-    ax.set_title("Mean Module Importance by Family and Cell Type", fontsize=11)
-    ax.legend(fontsize=8, loc="upper right")
-    ax.axhline(0, color="gray", linewidth=0.5, linestyle=":")
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig4_family_mean_importance")
-    print("  Saved fig4_family_mean_importance")
-
-    # Save data
-    summary_df = pd.DataFrame(mean_by_family, index=label_classes, columns=FAMILY_ORDER)
-    summary_df.to_csv(output_dir / "family_mean_importance.csv")
-    print("  Saved family_mean_importance.csv")
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Family-level module importance analysis")
-    p.add_argument("--results_dir", required=True,
-                   help="Importance results directory (with importance_matrix.npy)")
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--top_k", type=int, nargs="+", default=[10, 20, 50],
-                   help="Top-K values for contribution analysis")
-    p.add_argument("--enrichment_k", type=int, default=50,
-                   help="K for enrichment test")
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    print(f"Loading results from {args.results_dir}...")
-    data = load_importance_results(Path(args.results_dir))
-    importance = data["importance"]
-    label_classes = data["label_classes"]
-    module_names = data["module_names"]
-
-    family_indices = compute_family_indices(module_names)
-    print(f"Module families: { {f: len(idx) for f, idx in family_indices.items()} }")
-    print(f"Cell types: {len(label_classes)}")
-    print()
-
-    print("[1/4] Importance distribution by family...")
-    plot_family_importance_distribution(importance, label_classes, module_names, output_dir)
-    print()
-
-    print("[2/4] Family contribution to top-K...")
-    plot_family_contribution_topk(importance, label_classes, module_names, output_dir,
-                                  k_values=args.top_k)
-    print()
-
-    print("[3/4] Family enrichment test...")
-    enrich_df = run_family_enrichment(importance, label_classes, module_names, output_dir,
-                                      K=args.enrichment_k)
-    sig_count = enrich_df["significant"].sum()
-    print(f"  Significant enrichments (Bonferroni p<0.05): {sig_count}/{len(enrich_df)}")
-    print()
-
-    print("[4/4] Mean importance summary...")
-    plot_family_summary(importance, label_classes, module_names, output_dir, K=args.enrichment_k)
-    print()
-
-    print("All analyses complete.")
-    print(f"Output directory: {output_dir}")
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/marker_gene_verification.py

@@ -1,321 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Experiment ③-C: Canonical Marker Gene Verification
-
-Verifies that known cell-type marker genes reside in high-importance modules.
-Runs entirely from saved results (CPU only, no model inference needed).
-
-Outputs:
-  - marker_ranks.csv: per-marker rank information
-  - marker_metrics.json: summary statistics
-  - fig_marker_ranks.pdf/png: publication-quality dot plot
-"""
-
-import argparse
-import json
-import sys
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-from scipy import stats
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-sys.path.insert(0, str(SCRIPT_DIR))
-
-from pub_utils import (
-    CANONICAL_MARKERS,
-    CELL_TYPE_COLORS,
-    find_gene_biomods,
-    get_biomod_family,
-    importance_rank_percentile,
-    load_importance_results,
-    load_biomod_gene_mapping,
-    savefig,
-    setup_publication_style,
-)
-
-
-def compute_marker_ranks(
-    importance: np.ndarray,
-    label_classes: list,
-    module_names: list,
-    module_gene_symbols: dict,
-) -> pd.DataFrame:
-    """For each cell type's canonical markers, find their importance rank."""
-    rows = []
-    module_name_to_idx = {n: i for i, n in enumerate(module_names)}
-    num_modules = len(module_names)
-
-    for cls_idx, cls_name in enumerate(label_classes):
-        markers = CANONICAL_MARKERS.get(cls_name, [])
-        imp_row = importance[cls_idx]
-        ranked_indices = np.argsort(-imp_row)
-
-        for marker in markers:
-            containing_modules = find_gene_biomods(marker, module_gene_symbols)
-            if not containing_modules:
-                rows.append({
-                    "cell_type": cls_name,
-                    "marker_gene": marker,
-                    "found_in_modules": False,
-                    "num_modules_containing": 0,
-                    "best_module": None,
-                    "best_module_family": None,
-                    "best_rank": None,
-                    "rank_percentile": None,
-                    "best_importance": None,
-                })
-                continue
-
-            module_indices = [
-                module_name_to_idx[p] for p in containing_modules
-                if p in module_name_to_idx
-            ]
-            if not module_indices:
-                continue
-
-            ranks = []
-            for pi in module_indices:
-                rank = int(np.where(ranked_indices == pi)[0][0])
-                ranks.append((rank, pi))
-            best_rank, best_pi = min(ranks, key=lambda x: x[0])
-
-            rows.append({
-                "cell_type": cls_name,
-                "marker_gene": marker,
-                "found_in_modules": True,
-                "num_modules_containing": len(containing_modules),
-                "best_module": module_names[best_pi],
-                "best_module_family": get_biomod_family(module_names[best_pi]),
-                "best_rank": best_rank + 1,
-                "rank_percentile": 100.0 * best_rank / num_modules,
-                "best_importance": float(imp_row[best_pi]),
-            })
-
-    return pd.DataFrame(rows)
-
-
-def compute_enrichment_statistics(
-    marker_df: pd.DataFrame, num_modules: int
-) -> dict:
-    """Test whether markers are enriched in top modules vs random expectation."""
-    found = marker_df[marker_df["found_in_modules"]].copy()
-    if found.empty:
-        return {}
-
-    percentiles = found["rank_percentile"].values
-
-    top_5_count = (percentiles < 5).sum()
-    top_10_count = (percentiles < 10).sum()
-    top_20_count = (percentiles < 20).sum()
-    total = len(percentiles)
-
-    binom_p5 = 1 - stats.binom.cdf(top_5_count - 1, total, 0.05)
-    binom_p10 = 1 - stats.binom.cdf(top_10_count - 1, total, 0.10)
-    binom_p20 = 1 - stats.binom.cdf(top_20_count - 1, total, 0.20)
-
-    median_percentile = float(np.median(percentiles))
-    wilcoxon_stat, wilcoxon_p = stats.wilcoxon(
-        percentiles - 50.0, alternative="less"
-    )
-
-    return {
-        "total_markers_tested": total,
-        "markers_not_found": int((~marker_df["found_in_modules"]).sum()),
-        "median_rank_percentile": median_percentile,
-        "mean_rank_percentile": float(np.mean(percentiles)),
-        "in_top_5pct": int(top_5_count),
-        "in_top_10pct": int(top_10_count),
-        "in_top_20pct": int(top_20_count),
-        "binom_pvalue_top5": float(binom_p5),
-        "binom_pvalue_top10": float(binom_p10),
-        "binom_pvalue_top20": float(binom_p20),
-        "wilcoxon_stat": float(wilcoxon_stat),
-        "wilcoxon_pvalue": float(wilcoxon_p),
-    }
-
-
-def plot_marker_ranks(marker_df: pd.DataFrame, output_path):
-    """Publication-quality table figure of marker gene rank percentiles."""
-    setup_publication_style()
-
-    found = marker_df[marker_df["found_in_modules"]].copy()
-    if found.empty:
-        return
-
-    cell_types = list(CANONICAL_MARKERS.keys())
-    cell_types = [ct for ct in cell_types if ct in found["cell_type"].values]
-
-    # Build table data: each row = cell type, columns = marker genes with ranks
-    table_rows = []
-    for ct in cell_types:
-        ct_data = found[found["cell_type"] == ct].sort_values("rank_percentile")
-        markers_str = []
-        for _, row in ct_data.iterrows():
-            rank = int(row["best_rank"])
-            pct = row["rank_percentile"]
-            markers_str.append(f"{row['marker_gene']} (#{rank}, {pct:.1f}%)")
-        table_rows.append(markers_str)
-
-    # Determine max markers per cell type for column count
-    max_markers = max(len(r) for r in table_rows)
-    col_labels = [f"Marker {i+1}" for i in range(min(max_markers, 6))]
-
-    # Truncate to 6 columns max for readability
-    table_data = []
-    for row in table_rows:
-        table_data.append(row[:6] + [""] * (6 - len(row[:6])))
-
-    fig, ax = plt.subplots(figsize=(12, 0.5 * len(cell_types) + 2))
-    ax.axis("off")
-
-    table = ax.table(
-        cellText=table_data,
-        rowLabels=cell_types,
-        colLabels=col_labels,
-        cellLoc="center",
-        rowLoc="right",
-        loc="center",
-    )
-    table.auto_set_font_size(False)
-    table.set_fontsize(7.5)
-    table.scale(1.0, 1.4)
-
-    # Color cells by rank percentile
-    for i, row in enumerate(table_data):
-        for j, cell_text in enumerate(row):
-            cell = table[i + 1, j]
-            if cell_text and "%" in cell_text:
-                pct = float(cell_text.split(",")[-1].strip().rstrip("%)"))
-                if pct < 1:
-                    cell.set_facecolor("#fee0d2")
-                elif pct < 5:
-                    cell.set_facecolor("#fff5eb")
-
-    ax.set_title(
-        "Canonical Marker Genes: Importance Rank (format: Gene (#rank, percentile))",
-        fontsize=10, pad=20,
-    )
-
-    fig.tight_layout()
-    savefig(fig, output_path)
-
-
-def plot_marker_summary_bar(marker_df: pd.DataFrame, output_path):
-    """Bar chart showing fraction of markers in top-k% per cell type."""
-    setup_publication_style()
-
-    found = marker_df[marker_df["found_in_modules"]].copy()
-    if found.empty:
-        return
-
-    cell_types = [ct for ct in CANONICAL_MARKERS if ct in found["cell_type"].values]
-    thresholds = [1, 2, 3, 5]
-    colors = ["#a50f15", "#E41A1C", "#FF7F00", "#4DAF4A"]
-
-    fig, ax = plt.subplots(figsize=(10, 4.5))
-    x = np.arange(len(cell_types))
-    width = 0.2
-
-    for j, (thresh, color) in enumerate(zip(thresholds, colors)):
-        fracs = []
-        for ct in cell_types:
-            ct_data = found[found["cell_type"] == ct]
-            n_in = (ct_data["rank_percentile"] < thresh).sum()
-            fracs.append(n_in / len(ct_data) if len(ct_data) > 0 else 0)
-        offset = (j - 1.5) * width
-        ax.bar(x + offset, fracs, width, color=color, alpha=0.85,
-               label=f"Top {thresh}%")
-
-    # Random expectation lines
-    for thresh in thresholds:
-        ax.axhline(thresh / 100.0, color="gray", linestyle=":",
-                   linewidth=0.5, alpha=0.5)
-
-    ax.set_xticks(x)
-    ax.set_xticklabels(cell_types, rotation=45, ha="right")
-    ax.set_ylabel("Fraction of markers in top-k%")
-    ax.set_title("Marker Gene Enrichment in Top Importance Modules")
-    ax.legend(loc="upper right", fontsize=8)
-    ax.set_ylim(0, 1.05)
-
-    fig.tight_layout()
-    savefig(fig, output_path)
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Marker gene verification")
-    p.add_argument("--importance_matrix", required=True,
-                   help="Path to importance_matrix.npy")
-    p.add_argument("--importance_csv", required=True,
-                   help="Path to importance_matrix.csv (for labels)")
-    p.add_argument("--biomod_data_dir", "--module_data_dir", required=True,
-                   help="Directory containing merged_biomods.json", dest='module_data_dir')
-    p.add_argument("--output_dir", required=True)
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    importance = np.load(args.importance_matrix)
-    imp_csv = pd.read_csv(args.importance_csv, index_col=0)
-    label_classes = list(imp_csv.index)
-    module_names = list(imp_csv.columns)
-
-    _, module_gene_symbols = load_biomod_gene_mapping(args.module_data_dir)
-
-    print(f"Importance matrix: {importance.shape}")
-    print(f"Cell types: {len(label_classes)}")
-    print(f"Modules: {len(module_names)}")
-    print(f"Canonical markers defined for: {len(CANONICAL_MARKERS)} cell types")
-
-    marker_df = compute_marker_ranks(
-        importance, label_classes, module_names, module_gene_symbols
-    )
-    marker_df.to_csv(output_dir / "marker_ranks.csv", index=False)
-
-    stats_dict = compute_enrichment_statistics(marker_df, len(module_names))
-
-    per_ct_stats = {}
-    for ct in label_classes:
-        ct_found = marker_df[
-            (marker_df["cell_type"] == ct) & marker_df["found_in_modules"]
-        ]
-        if len(ct_found) > 0:
-            per_ct_stats[ct] = {
-                "n_markers": len(ct_found),
-                "median_percentile": float(ct_found["rank_percentile"].median()),
-                "best_marker": ct_found.loc[
-                    ct_found["rank_percentile"].idxmin(), "marker_gene"
-                ],
-                "best_rank": int(ct_found["best_rank"].min()),
-            }
-
-    metrics = {**stats_dict, "per_cell_type": per_ct_stats}
-    with open(output_dir / "marker_metrics.json", "w") as f:
-        json.dump(metrics, f, indent=2)
-
-    print(f"\nResults:")
-    print(f"  Markers found in modules: {stats_dict['total_markers_tested']}")
-    print(f"  Markers NOT found: {stats_dict['markers_not_found']}")
-    print(f"  Median rank percentile: {stats_dict['median_rank_percentile']:.1f}%")
-    print(f"  In top 5%: {stats_dict['in_top_5pct']}")
-    print(f"  In top 10%: {stats_dict['in_top_10pct']}")
-    print(f"  In top 20%: {stats_dict['in_top_20pct']}")
-    print(f"  Binomial p-value (top 10%): {stats_dict['binom_pvalue_top10']:.2e}")
-    print(f"  Wilcoxon p-value (< 50%): {stats_dict['wilcoxon_pvalue']:.2e}")
-
-    plot_marker_ranks(marker_df, output_dir / "fig_marker_ranks")
-    plot_marker_summary_bar(marker_df, output_dir / "fig_marker_summary_bar")
-
-    print(f"\nFigures saved to {output_dir}")
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/masking_ablation.py

@@ -1,375 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Experiment ③-A: Masking Ablation Curves
-
-Progressively mask top-k / bottom-k / random-k modules and measure accuracy
-drop. Proves that top-importance modules are necessary for classification.
-
-Requires GPU (model inference). Uses post_encoder masking for efficiency.
-
-Outputs:
-  - ablation_curves.npz: accuracy arrays
-  - ablation_metrics.json: summary statistics
-  - per_cell_deltas.npy: per-cell importance deltas (for permutation tests)
-  - attention_weights.npy: GatedAttentionPool weights (bonus)
-"""
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-from typing import List, Tuple
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-INTERP_DIR = SCRIPT_DIR.parent
-BIOMOD_ROOT = INTERP_DIR.parent
-SCLONG_ROOT = BIOMOD_ROOT.parent
-
-for p in (str(SCRIPT_DIR), str(BIOMOD_ROOT), str(SCLONG_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from biomod_importance import (
-    CellDataset,
-    _classify_from_tokens,
-    load_finetuned_model,
-    load_test_data,
-)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger(__name__)
-
-
-@torch.no_grad()
-def encode_all_cells(model, dataloader, preprocessor, reindex_state, device):
-    """Encode all test cells once, return module tokens and labels."""
-    all_tokens = []
-    all_labels = []
-    for x_batch, y_batch in dataloader:
-        x_batch = x_batch.to(device)
-        x_model = preprocessor.reindex_tensor_to_model_space(
-            x_batch, reindex_state=reindex_state
-        )
-        latent, _, ids_restore, _ = model.backbone.forward_encoder(
-            x_model, mask_ratio=0.0
-        )
-        if model.head_type == "cls":
-            module_tokens = latent[:, 1:, :]
-        else:
-            module_tokens = model._restore_biomod_order(latent, ids_restore)
-        all_tokens.append(module_tokens.cpu())
-        all_labels.append(y_batch.numpy())
-
-    return torch.cat(all_tokens, dim=0), np.concatenate(all_labels)
-
-
-@torch.no_grad()
-def compute_per_cell_deltas(model, all_tokens, all_labels, device, batch_size=256):
-    """Compute per-cell, per-module importance deltas.
-
-    Returns: np.ndarray of shape (n_cells, n_modules)
-    """
-    n_cells, n_modules, _ = all_tokens.shape
-    deltas = np.zeros((n_cells, n_modules), dtype=np.float32)
-
-    for start in range(0, n_cells, batch_size):
-        end = min(start + batch_size, n_cells)
-        tokens = all_tokens[start:end].to(device)
-        labels = torch.from_numpy(all_labels[start:end]).to(device)
-        batch_idx = torch.arange(end - start, device=device)
-
-        baseline_logits = _classify_from_tokens(model, tokens)
-        baseline_correct = baseline_logits[batch_idx, labels].cpu().numpy()
-
-        for pi in range(n_modules):
-            masked = tokens.clone()
-            masked[:, pi, :] = 0.0
-            masked_logits = _classify_from_tokens(model, masked)
-            masked_correct = masked_logits[batch_idx, labels].cpu().numpy()
-            deltas[start:end, pi] = baseline_correct - masked_correct
-
-    return deltas
-
-
-@torch.no_grad()
-def extract_attention_weights(model, all_tokens, device, batch_size=256):
-    """Extract GatedAttentionPool weights for all cells.
-
-    Returns: np.ndarray of shape (n_cells, n_modules) or None if not attn_pool.
-    """
-    if model.head_type != "attn_pool":
-        return None
-
-    n_cells = all_tokens.shape[0]
-    all_weights = []
-
-    for start in range(0, n_cells, batch_size):
-        end = min(start + batch_size, n_cells)
-        tokens = all_tokens[start:end].to(device)
-        _, weights = model.pooler(tokens)
-        all_weights.append(weights.cpu().numpy())
-
-    return np.concatenate(all_weights, axis=0)
-
-
-@torch.no_grad()
-def compute_masked_accuracy_post_encoder(
-    model, all_tokens, all_labels, mask_module_indices, device, batch_size=256
-):
-    """Compute accuracy with specified modules zeroed out (post-encoder)."""
-    n_cells = all_tokens.shape[0]
-    correct = 0
-
-    for start in range(0, n_cells, batch_size):
-        end = min(start + batch_size, n_cells)
-        tokens = all_tokens[start:end].to(device).clone()
-        labels = all_labels[start:end]
-
-        for pi in mask_module_indices:
-            tokens[:, pi, :] = 0.0
-
-        logits = _classify_from_tokens(model, tokens)
-        preds = logits.argmax(dim=-1).cpu().numpy()
-        correct += (preds == labels).sum()
-
-    return correct / n_cells
-
-
-def run_ablation_curves(
-    model, all_tokens, all_labels, importance_matrix,
-    label_classes, k_values, n_random_trials, device, batch_size=256,
-):
-    """Run ablation for all cell types and masking orders."""
-    num_classes = len(label_classes)
-    num_modules = all_tokens.shape[1]
-    n_k = len(k_values)
-
-    # Overall accuracy (all cell types together)
-    overall_top = np.zeros(n_k)
-    overall_bottom = np.zeros(n_k)
-    overall_random_mean = np.zeros(n_k)
-    overall_random_std = np.zeros(n_k)
-
-    # Per-class accuracy
-    per_class_top = np.zeros((num_classes, n_k))
-    per_class_bottom = np.zeros((num_classes, n_k))
-    per_class_random_mean = np.zeros((num_classes, n_k))
-    per_class_random_std = np.zeros((num_classes, n_k))
-
-    # --- Overall ablation (use mean importance across classes) ---
-    mean_importance = importance_matrix.mean(axis=0)
-    overall_ranking = np.argsort(-mean_importance)
-
-    logger.info("Computing overall ablation curves...")
-    for ki, k in enumerate(k_values):
-        top_indices = overall_ranking[:k].tolist()
-        bottom_indices = overall_ranking[-k:].tolist()
-
-        overall_top[ki] = compute_masked_accuracy_post_encoder(
-            model, all_tokens, all_labels, top_indices, device, batch_size
-        )
-        overall_bottom[ki] = compute_masked_accuracy_post_encoder(
-            model, all_tokens, all_labels, bottom_indices, device, batch_size
-        )
-
-        random_accs = []
-        rng = np.random.default_rng(2021)
-        for _ in range(n_random_trials):
-            rand_indices = rng.choice(num_modules, k, replace=False).tolist()
-            acc = compute_masked_accuracy_post_encoder(
-                model, all_tokens, all_labels, rand_indices, device, batch_size
-            )
-            random_accs.append(acc)
-        overall_random_mean[ki] = np.mean(random_accs)
-        overall_random_std[ki] = np.std(random_accs)
-
-        logger.info(
-            "  k=%d: top=%.4f, bottom=%.4f, random=%.4f±%.4f",
-            k, overall_top[ki], overall_bottom[ki],
-            overall_random_mean[ki], overall_random_std[ki],
-        )
-
-    # --- Per-class ablation ---
-    for cls_idx in range(num_classes):
-        cls_mask = all_labels == cls_idx
-        if cls_mask.sum() == 0:
-            continue
-        cls_tokens = all_tokens[cls_mask]
-        cls_labels = all_labels[cls_mask]
-        cls_ranking = np.argsort(-importance_matrix[cls_idx])
-
-        logger.info("Computing ablation for %s (%d cells)...",
-                    label_classes[cls_idx], cls_mask.sum())
-
-        for ki, k in enumerate(k_values):
-            top_indices = cls_ranking[:k].tolist()
-            bottom_indices = cls_ranking[-k:].tolist()
-
-            per_class_top[cls_idx, ki] = compute_masked_accuracy_post_encoder(
-                model, cls_tokens, cls_labels, top_indices, device, batch_size
-            )
-            per_class_bottom[cls_idx, ki] = compute_masked_accuracy_post_encoder(
-                model, cls_tokens, cls_labels, bottom_indices, device, batch_size
-            )
-
-            random_accs = []
-            rng = np.random.default_rng(2021 + cls_idx)
-            for _ in range(n_random_trials):
-                rand_indices = rng.choice(num_modules, k, replace=False).tolist()
-                acc = compute_masked_accuracy_post_encoder(
-                    model, cls_tokens, cls_labels, rand_indices, device, batch_size
-                )
-                random_accs.append(acc)
-            per_class_random_mean[cls_idx, ki] = np.mean(random_accs)
-            per_class_random_std[cls_idx, ki] = np.std(random_accs)
-
-    return {
-        "k_values": np.array(k_values),
-        "overall_top": overall_top,
-        "overall_bottom": overall_bottom,
-        "overall_random_mean": overall_random_mean,
-        "overall_random_std": overall_random_std,
-        "per_class_top": per_class_top,
-        "per_class_bottom": per_class_bottom,
-        "per_class_random_mean": per_class_random_mean,
-        "per_class_random_std": per_class_random_std,
-        "label_classes": label_classes,
-    }
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Masking ablation curves")
-    p.add_argument("--finetune_checkpoint", required=True)
-    p.add_argument("--config_pkl", required=True)
-    p.add_argument("--biomod_data_dir", "--module_data_dir", required=True, dest='module_data_dir')
-    p.add_argument("--gene_list_path", required=True)
-    p.add_argument("--gene_mapping_path", required=True)
-    p.add_argument("--gene2vec_data_path", required=True)
-    p.add_argument("--gene2ens_file", required=True)
-    p.add_argument("--data_path", required=True)
-    p.add_argument("--label_column", default="celltype")
-    p.add_argument("--head_type", default="attn_pool",
-                   help="Classification head type (attn_pool, module_mean, etc.)")
-    p.add_argument("--gene_name_column", default="auto",
-                   help="Gene name column in h5ad (auto=detect)")
-    p.add_argument("--importance_matrix", required=True,
-                   help="Path to pre-computed importance_matrix.npy")
-    p.add_argument("--k_values", type=int, nargs="+",
-                   default=[1, 2, 5, 10, 20, 50, 100, 200])
-    p.add_argument("--n_random_trials", type=int, default=10)
-    p.add_argument("--batch_size", type=int, default=256)
-    p.add_argument("--test_ratio", type=float, default=0.2)
-    p.add_argument("--seed", type=int, default=2021)
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--save_deltas", action="store_true",
-                   help="Save per-cell delta matrix (for permutation tests)")
-    p.add_argument("--device", default="cuda")
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-    logger.info("Device: %s", device)
-
-    for attr in ("finetune_checkpoint", "config_pkl", "module_data_dir",
-                 "gene_list_path", "gene_mapping_path", "gene2vec_data_path",
-                 "gene2ens_file", "data_path", "importance_matrix"):
-        val = getattr(args, attr, None)
-        if val is not None and not Path(val).is_absolute():
-            setattr(args, attr, str(Path(val).resolve()))
-
-    logger.info("Loading model...")
-    model, biomod_data, label_classes, head_type = load_finetuned_model(args, device)
-    num_modules = len(biomod_data)
-
-    logger.info("Loading test data...")
-    X_test, y_test, le, preprocessor, reindex_state = load_test_data(
-        args, label_classes, device
-    )
-    dataset = CellDataset(X_test, y_test)
-    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
-
-    logger.info("Loading importance matrix...")
-    importance_matrix = np.load(args.importance_matrix)
-
-    logger.info("Encoding all cells (encoder runs once)...")
-    all_tokens, all_labels = encode_all_cells(
-        model, dataloader, preprocessor, reindex_state, device
-    )
-    logger.info("Encoded %d cells, %d modules", all_tokens.shape[0], all_tokens.shape[1])
-
-    # Baseline accuracy
-    baseline_acc = compute_masked_accuracy_post_encoder(
-        model, all_tokens, all_labels, [], device, args.batch_size
-    )
-    logger.info("Baseline accuracy: %.4f", baseline_acc)
-
-    # Extract attention weights (bonus)
-    logger.info("Extracting attention weights...")
-    attn_weights = extract_attention_weights(model, all_tokens, device, args.batch_size)
-    if attn_weights is not None:
-        np.save(output_dir / "attention_weights.npy", attn_weights)
-        logger.info("Saved attention weights: %s", attn_weights.shape)
-
-    # Save per-cell deltas if requested
-    if args.save_deltas:
-        logger.info("Computing per-cell deltas (for permutation tests)...")
-        deltas = compute_per_cell_deltas(
-            model, all_tokens, all_labels, device, args.batch_size
-        )
-        np.save(output_dir / "per_cell_deltas.npy", deltas)
-        np.save(output_dir / "per_cell_labels.npy", all_labels)
-        logger.info("Saved per-cell deltas: %s", deltas.shape)
-
-    # Run ablation curves
-    logger.info("Running ablation curves...")
-    results = run_ablation_curves(
-        model, all_tokens, all_labels, importance_matrix,
-        label_classes, args.k_values, args.n_random_trials,
-        device, args.batch_size,
-    )
-
-    # Save results
-    np.savez(
-        output_dir / "ablation_curves.npz",
-        k_values=results["k_values"],
-        overall_top=results["overall_top"],
-        overall_bottom=results["overall_bottom"],
-        overall_random_mean=results["overall_random_mean"],
-        overall_random_std=results["overall_random_std"],
-        per_class_top=results["per_class_top"],
-        per_class_bottom=results["per_class_bottom"],
-        per_class_random_mean=results["per_class_random_mean"],
-        per_class_random_std=results["per_class_random_std"],
-    )
-
-    metrics = {
-        "baseline_accuracy": baseline_acc,
-        "k_values": args.k_values,
-        "n_random_trials": args.n_random_trials,
-        "num_test_cells": int(all_tokens.shape[0]),
-        "num_modules": num_modules,
-        "label_classes": label_classes,
-        "overall_top_acc": results["overall_top"].tolist(),
-        "overall_bottom_acc": results["overall_bottom"].tolist(),
-        "overall_random_mean_acc": results["overall_random_mean"].tolist(),
-    }
-    with open(output_dir / "ablation_metrics.json", "w") as f:
-        json.dump(metrics, f, indent=2)
-
-    logger.info("Done. Results saved to %s", output_dir)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/pub_utils.py

@@ -1,206 +0,0 @@
-"""Shared utilities for publication-quality module importance analysis."""
-
-import json
-from pathlib import Path
-from typing import Dict, List, Tuple
-
-import matplotlib as mpl
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-
-# ─── Canonical Marker Genes ─────────────────────────────────────────────────
-# Well-established markers for Zheng68K cell types (from literature)
-
-CANONICAL_MARKERS: Dict[str, List[str]] = {
-    "CD14+ Monocyte": [
-        "CD14", "LYZ", "S100A8", "S100A9", "VCAN", "FCN1", "CD68",
-    ],
-    "CD19+ B": [
-        "CD19", "MS4A1", "CD79A", "CD79B", "PAX5", "BANK1",
-    ],
-    "CD34+": [
-        "CD34", "KIT", "THY1", "CRHBP", "SPINK2",
-    ],
-    "CD4+ T Helper2": [
-        "CD4", "GATA3", "IL7R", "MAL", "CCR4",
-    ],
-    "CD4+/CD25 T Reg": [
-        "CD4", "FOXP3", "IL2RA", "CTLA4", "TIGIT", "IKZF2",
-    ],
-    "CD4+/CD45RA+/CD25- Naive T": [
-        "CD4", "CCR7", "SELL", "LEF1", "TCF7",
-    ],
-    "CD4+/CD45RO+ Memory": [
-        "CD4", "IL7R", "S100A4", "ANXA1", "GPR183",
-    ],
-    "CD56+ NK": [
-        "NCAM1", "NKG7", "GNLY", "KLRD1", "KLRF1", "GZMB",
-    ],
-    "CD8+ Cytotoxic T": [
-        "CD8A", "CD8B", "GZMB", "PRF1", "IFNG", "NKG7",
-    ],
-    "CD8+/CD45RA+ Naive Cytotoxic": [
-        "CD8A", "CCR7", "SELL", "LEF1", "TCF7",
-    ],
-    "Dendritic": [
-        "FCER1A", "CD1C", "CLEC10A", "HLA-DRA", "ITGAX",
-    ],
-}
-
-# ─── Color Palettes ─────────────────────────────────────────────────────────
-
-CELL_TYPE_COLORS: Dict[str, str] = {
-    "CD14+ Monocyte": "#E41A1C",
-    "CD19+ B": "#377EB8",
-    "CD34+": "#4DAF4A",
-    "CD4+ T Helper2": "#984EA3",
-    "CD4+/CD25 T Reg": "#FF7F00",
-    "CD4+/CD45RA+/CD25- Naive T": "#A65628",
-    "CD4+/CD45RO+ Memory": "#F781BF",
-    "CD56+ NK": "#66C2A5",
-    "CD8+ Cytotoxic T": "#FC8D62",
-    "CD8+/CD45RA+ Naive Cytotoxic": "#8DA0CB",
-    "Dendritic": "#E78AC3",
-}
-
-MODULE_FAMILY_COLORS: Dict[str, str] = {
-    "Reactome": "#1B9E77",
-    "Hallmarks": "#D95F02",
-    "Hotspot": "#7570B3",
-    "Processed": "#E7298A",
-}
-
-
-# ─── Publication Style ───────────────────────────────────────────────────────
-
-def setup_publication_style():
-    """Set matplotlib rcParams for publication-quality figures."""
-    mpl.rcParams.update({
-        "figure.dpi": 300,
-        "savefig.dpi": 300,
-        "font.family": "sans-serif",
-        "font.sans-serif": ["Arial", "DejaVu Sans", "Helvetica"],
-        "font.size": 9,
-        "axes.titlesize": 11,
-        "axes.labelsize": 10,
-        "xtick.labelsize": 8,
-        "ytick.labelsize": 8,
-        "legend.fontsize": 8,
-        "figure.titlesize": 12,
-        "lines.linewidth": 1.5,
-        "lines.markersize": 4,
-        "axes.linewidth": 0.8,
-        "xtick.major.width": 0.8,
-        "ytick.major.width": 0.8,
-        "axes.spines.top": False,
-        "axes.spines.right": False,
-        "pdf.fonttype": 42,
-        "ps.fonttype": 42,
-    })
-
-
-def savefig(fig, path_stem, dpi=300, tight=True):
-    """Save figure as both PDF and PNG."""
-    path = Path(path_stem)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    kwargs = {"dpi": dpi, "bbox_inches": "tight"} if tight else {"dpi": dpi}
-    fig.savefig(path.with_suffix(".pdf"), **kwargs)
-    fig.savefig(path.with_suffix(".png"), **kwargs)
-    plt.close(fig)
-
-
-# ─── Data Loading ────────────────────────────────────────────────────────────
-
-def load_importance_results(results_dir) -> dict:
-    """Load importance matrix and associated metadata from a results directory."""
-    results_dir = Path(results_dir)
-    data = {}
-    data["importance"] = np.load(results_dir / "importance_matrix.npy")
-    data["importance_csv"] = pd.read_csv(
-        results_dir / "importance_matrix.csv", index_col=0
-    )
-    data["label_classes"] = list(data["importance_csv"].index)
-    data["module_names"] = list(data["importance_csv"].columns)
-
-    metrics_path = results_dir / "metrics.json"
-    if metrics_path.exists():
-        with open(metrics_path) as f:
-            data["metrics"] = json.load(f)
-
-    top_modules_path = results_dir / "top_modules_per_celltype.json"
-    if top_modules_path.exists():
-        with open(top_modules_path) as f:
-            data["top_modules"] = json.load(f)
-
-    enrichment_path = results_dir / "enrichment_results.csv"
-    if enrichment_path.exists():
-        data["enrichment"] = pd.read_csv(enrichment_path)
-
-    return data
-
-
-def load_biomod_gene_mapping(module_data_dir) -> Tuple[List[str], Dict[str, List[str]]]:
-    """Load merged_biomods.json and return (module_names, module_gene_symbols).
-
-    Returns:
-        module_names: ordered list of module names
-        module_gene_symbols: dict mapping module_name -> list of gene symbols
-    """
-    module_data_dir = Path(module_data_dir)
-    with open(module_data_dir / "merged_biomods.json") as f:
-        modules = json.load(f)
-    module_names = list(modules.keys())
-    module_gene_symbols = {}
-    for name, genes in modules.items():
-        if genes and isinstance(genes[0], str):
-            module_gene_symbols[name] = genes
-        else:
-            module_gene_symbols[name] = [str(g) for g in genes]
-    return module_names, module_gene_symbols
-
-
-def get_biomod_family(module_name: str) -> str:
-    """Determine which family a module belongs to based on its name prefix."""
-    if module_name.startswith("module1_hallmarks"):
-        return "Hallmarks"
-    elif module_name.startswith("module2_reactome"):
-        return "Reactome"
-    elif module_name.startswith("module3_hotspot"):
-        return "Hotspot"
-    elif module_name.startswith("module3_processed"):
-        return "Processed"
-    return "Unknown"
-
-
-def get_biomod_families(module_names: List[str]) -> List[str]:
-    """Get family labels for a list of module names."""
-    return [get_biomod_family(p) for p in module_names]
-
-
-# ─── Analysis Helpers ────────────────────────────────────────────────────────
-
-def find_gene_biomods(
-    gene_symbol: str, module_gene_symbols: Dict[str, List[str]]
-) -> List[str]:
-    """Find all modules containing a given gene symbol."""
-    gene_upper = gene_symbol.upper()
-    return [
-        name for name, genes in module_gene_symbols.items()
-        if gene_upper in [g.upper() for g in genes]
-    ]
-
-
-def importance_rank_percentile(
-    importance_row: np.ndarray, module_indices: List[int]
-) -> float:
-    """Compute the best (lowest) rank percentile for given module indices.
-
-    Returns percentile in [0, 100] where 0 = most important.
-    """
-    if not module_indices:
-        return 100.0
-    ranked = np.argsort(-importance_row)
-    ranks = np.array([np.where(ranked == pi)[0][0] for pi in module_indices])
-    best_rank = ranks.min()
-    return 100.0 * best_rank / len(importance_row)

BioMod/interpretability/biomod_importance/random_baseline.py

@@ -1,333 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Experiment ③-B: Random Baseline / Permutation Tests
-
-Two modes:
-  A) shuffle_modules (GPU): Permute module-gene assignments, recompute importance.
-     Proves importance is not random.
-  B) shuffle_labels (CPU): Permute cell type labels, re-aggregate per-cell deltas.
-     Proves importance is cell-type-specific.
-
-Outputs:
-  - permutation_results.npz: null distributions
-  - permutation_metrics.json: p-values and effect sizes
-"""
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-INTERP_DIR = SCRIPT_DIR.parent
-BIOMOD_ROOT = INTERP_DIR.parent
-SCLONG_ROOT = BIOMOD_ROOT.parent
-
-for p in (str(SCRIPT_DIR), str(BIOMOD_ROOT), str(SCLONG_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from biomod_importance import (
-    CellDataset,
-    _classify_from_tokens,
-    load_finetuned_model,
-    load_test_data,
-)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger(__name__)
-
-
-# ─── Mode A: Shuffle Modules (GPU) ──────────────────────────────────────────
-
-@torch.no_grad()
-def compute_importance_shuffled_biomods(
-    model, dataloader, preprocessor, reindex_state,
-    num_modules, num_classes, device, seed,
-):
-    """Compute importance with shuffled module-gene assignments.
-
-    Strategy: After encoding, randomly permute the module token order.
-    This is equivalent to shuffling which module each token represents,
-    breaking the learned module-gene correspondence.
-    """
-    rng = np.random.default_rng(seed)
-    perm = rng.permutation(num_modules)
-
-    importance = np.zeros((num_classes, num_modules), dtype=np.float64)
-    class_counts = np.zeros(num_classes, dtype=np.float64)
-
-    for x_batch, y_batch in dataloader:
-        x_batch = x_batch.to(device)
-        y_batch = y_batch.numpy()
-        x_model = preprocessor.reindex_tensor_to_model_space(
-            x_batch, reindex_state=reindex_state
-        )
-
-        latent, _, ids_restore, _ = model.backbone.forward_encoder(
-            x_model, mask_ratio=0.0
-        )
-        if model.head_type == "cls":
-            module_tokens = latent[:, 1:, :]
-        else:
-            module_tokens = model._restore_biomod_order(latent, ids_restore)
-
-        # Permute module token order (breaks module identity)
-        module_tokens_shuffled = module_tokens[:, perm, :]
-
-        baseline_logits = _classify_from_tokens(model, module_tokens_shuffled)
-        y_tensor = torch.from_numpy(y_batch).to(device)
-        batch_idx = torch.arange(len(y_batch), device=device)
-        baseline_correct = baseline_logits[batch_idx, y_tensor].cpu().numpy()
-
-        for pi in range(num_modules):
-            masked = module_tokens_shuffled.clone()
-            masked[:, pi, :] = 0.0
-            masked_logits = _classify_from_tokens(model, masked)
-            masked_correct = masked_logits[batch_idx, y_tensor].cpu().numpy()
-
-            delta = baseline_correct - masked_correct
-            for cls_id in range(num_classes):
-                mask_cls = y_batch == cls_id
-                if mask_cls.any():
-                    importance[cls_id, pi] += delta[mask_cls].sum()
-
-        for cls_id in range(num_classes):
-            class_counts[cls_id] += (y_batch == cls_id).sum()
-
-    class_counts[class_counts == 0] = 1.0
-    importance /= class_counts[:, None]
-    return importance
-
-
-# ─── Mode B: Shuffle Labels (CPU) ───────────────────────────────────────────
-
-def compute_importance_shuffled_labels(
-    per_cell_deltas: np.ndarray,
-    labels: np.ndarray,
-    num_classes: int,
-    n_permutations: int,
-    seed: int,
-) -> np.ndarray:
-    """Re-aggregate per-cell deltas with shuffled labels.
-
-    Args:
-        per_cell_deltas: shape (n_cells, n_modules)
-        labels: shape (n_cells,) integer class labels
-        num_classes: number of classes
-        n_permutations: number of permutations
-        seed: random seed
-
-    Returns:
-        null_importance: shape (n_permutations, num_classes, n_modules)
-    """
-    n_cells, n_modules = per_cell_deltas.shape
-    rng = np.random.default_rng(seed)
-    null_importance = np.zeros(
-        (n_permutations, num_classes, n_modules), dtype=np.float32
-    )
-
-    for perm_idx in range(n_permutations):
-        shuffled_labels = rng.permutation(labels)
-        for cls_id in range(num_classes):
-            mask = shuffled_labels == cls_id
-            count = mask.sum()
-            if count > 0:
-                null_importance[perm_idx, cls_id] = (
-                    per_cell_deltas[mask].sum(axis=0) / count
-                )
-
-    return null_importance
-
-
-# ─── Statistics ──────────────────────────────────────────────────────────────
-
-def compute_permutation_pvalues(
-    real_importance: np.ndarray,
-    null_importance: np.ndarray,
-) -> dict:
-    """Compute empirical p-values from permutation null.
-
-    Args:
-        real_importance: shape (num_classes, num_modules)
-        null_importance: shape (n_perms, num_classes, num_modules)
-    """
-    n_perms = null_importance.shape[0]
-
-    # Per-cell-type max importance
-    real_max = real_importance.max(axis=1)  # (num_classes,)
-    null_max = null_importance.max(axis=2)  # (n_perms, num_classes)
-
-    pvalues_max = np.zeros(real_importance.shape[0])
-    for cls_id in range(real_importance.shape[0]):
-        pvalues_max[cls_id] = (
-            (null_max[:, cls_id] >= real_max[cls_id]).sum() + 1
-        ) / (n_perms + 1)
-
-    # Global max
-    real_global_max = real_importance.max()
-    null_global_max = null_importance.max(axis=(1, 2))
-    pvalue_global = ((null_global_max >= real_global_max).sum() + 1) / (n_perms + 1)
-
-    # Mean importance (top-10 per class)
-    real_top10_mean = np.sort(real_importance, axis=1)[:, -10:].mean(axis=1)
-    null_top10_mean = np.sort(null_importance, axis=2)[:, :, -10:].mean(axis=2)
-
-    pvalues_top10 = np.zeros(real_importance.shape[0])
-    for cls_id in range(real_importance.shape[0]):
-        pvalues_top10[cls_id] = (
-            (null_top10_mean[:, cls_id] >= real_top10_mean[cls_id]).sum() + 1
-        ) / (n_perms + 1)
-
-    return {
-        "pvalue_global_max": float(pvalue_global),
-        "pvalues_per_class_max": pvalues_max.tolist(),
-        "pvalues_per_class_top10_mean": pvalues_top10.tolist(),
-        "real_global_max": float(real_global_max),
-        "null_global_max_mean": float(null_global_max.mean()),
-        "null_global_max_std": float(null_global_max.std()),
-        "effect_size": float(
-            (real_global_max - null_global_max.mean()) / max(null_global_max.std(), 1e-8)
-        ),
-    }
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="Permutation tests for module importance")
-    p.add_argument("--mode", choices=["shuffle_modules", "shuffle_labels", "both"],
-                   default="both")
-    # Model args (for shuffle_modules mode)
-    p.add_argument("--finetune_checkpoint")
-    p.add_argument("--config_pkl")
-    p.add_argument("--biomod_data_dir", "--module_data_dir", dest='module_data_dir')
-    p.add_argument("--gene_list_path")
-    p.add_argument("--gene_mapping_path")
-    p.add_argument("--gene2vec_data_path")
-    p.add_argument("--gene2ens_file")
-    p.add_argument("--data_path")
-    p.add_argument("--head_type", default="attn_pool")
-    p.add_argument("--gene_name_column", default="auto",
-                   help="Gene name column in h5ad (auto=detect)")
-    p.add_argument("--label_column", default="celltype")
-    p.add_argument("--test_ratio", type=float, default=0.2)
-    # Shared args
-    p.add_argument("--importance_matrix", required=True,
-                   help="Path to real importance_matrix.npy")
-    p.add_argument("--per_cell_deltas",
-                   help="Path to per_cell_deltas.npy (for shuffle_labels mode)")
-    p.add_argument("--per_cell_labels",
-                   help="Path to per_cell_labels.npy")
-    p.add_argument("--n_permutations", type=int, default=100)
-    p.add_argument("--batch_size", type=int, default=256)
-    p.add_argument("--seed", type=int, default=2021)
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--device", default="cuda")
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    real_importance = np.load(args.importance_matrix)
-    num_classes, num_modules = real_importance.shape
-    logger.info("Real importance: %s", real_importance.shape)
-
-    results = {}
-
-    # ─── Mode A: Shuffle Modules ────────────────────────────────────────
-    if args.mode in ("shuffle_modules", "both"):
-        assert args.finetune_checkpoint, "--finetune_checkpoint required for shuffle_modules"
-        device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-
-        for attr in ("finetune_checkpoint", "config_pkl", "module_data_dir",
-                     "gene_list_path", "gene_mapping_path", "gene2vec_data_path",
-                     "gene2ens_file", "data_path"):
-            val = getattr(args, attr, None)
-            if val is not None and not Path(val).is_absolute():
-                setattr(args, attr, str(Path(val).resolve()))
-
-        logger.info("Loading model for shuffle_modules mode...")
-        model, biomod_data, label_classes, _ = load_finetuned_model(args, device)
-
-        logger.info("Loading test data...")
-        X_test, y_test, le, preprocessor, reindex_state = load_test_data(
-            args, label_classes, device
-        )
-        dataset = CellDataset(X_test, y_test)
-        dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
-
-        n_perms_a = min(args.n_permutations, 10)
-        logger.info("Running %d module-shuffle permutations...", n_perms_a)
-
-        null_module_shuffle = np.zeros(
-            (n_perms_a, num_classes, num_modules), dtype=np.float32
-        )
-        for pi in range(n_perms_a):
-            logger.info("  Permutation %d/%d...", pi + 1, n_perms_a)
-            null_module_shuffle[pi] = compute_importance_shuffled_biomods(
-                model, dataloader, preprocessor, reindex_state,
-                num_modules, num_classes, device,
-                seed=args.seed + pi + 1000,
-            )
-
-        np.save(output_dir / "null_module_shuffle.npy", null_module_shuffle)
-        stats_a = compute_permutation_pvalues(real_importance, null_module_shuffle)
-        stats_a["n_permutations"] = n_perms_a
-        results["shuffle_modules"] = stats_a
-        logger.info("Module-shuffle: global p=%.4f, effect_size=%.2f",
-                    stats_a["pvalue_global_max"], stats_a["effect_size"])
-
-    # ─── Mode B: Shuffle Labels ─────────────────────────────────────────
-    if args.mode in ("shuffle_labels", "both"):
-        if args.per_cell_deltas:
-            deltas_path = Path(args.per_cell_deltas)
-            labels_path = Path(args.per_cell_labels) if args.per_cell_labels else (
-                deltas_path.parent / "per_cell_labels.npy"
-            )
-            logger.info("Loading per-cell deltas from %s", deltas_path)
-            per_cell_deltas = np.load(deltas_path)
-            labels = np.load(labels_path)
-        else:
-            logger.error("--per_cell_deltas required for shuffle_labels mode")
-            logger.error("Run masking_ablation.py with --save_deltas first")
-            sys.exit(1)
-
-        n_perms_b = args.n_permutations
-        logger.info("Running %d label-shuffle permutations...", n_perms_b)
-
-        null_label_shuffle = compute_importance_shuffled_labels(
-            per_cell_deltas, labels, num_classes, n_perms_b, args.seed
-        )
-
-        np.save(output_dir / "null_label_shuffle.npy", null_label_shuffle)
-        stats_b = compute_permutation_pvalues(real_importance, null_label_shuffle)
-        stats_b["n_permutations"] = n_perms_b
-        results["shuffle_labels"] = stats_b
-        logger.info("Label-shuffle: global p=%.4f, effect_size=%.2f",
-                    stats_b["pvalue_global_max"], stats_b["effect_size"])
-
-    # ─── Save metrics (merge with existing to avoid overwriting) ────────
-    metrics_path = output_dir / "permutation_metrics.json"
-    if metrics_path.exists():
-        with open(metrics_path) as f:
-            existing = json.load(f)
-        existing.update(results)
-        results = existing
-
-    with open(metrics_path, "w") as f:
-        json.dump(results, f, indent=2)
-
-    logger.info("Done. Results saved to %s", output_dir)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_importance/tcell_reactome_analysis.py

@@ -1,466 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-T Cell Reactome Pathway Importance Analysis
-
-Analyzes which Reactome pathways are most important for T cell subtype
-classification, using pre-computed importance matrix (no GPU needed).
-
-Outputs:
-  1. Per T cell subtype top Reactome pathway ranking (table + bar chart)
-  2. Shared vs unique critical pathways across T cell subtypes
-  3. T cell-focused Reactome importance heatmap
-  4. Summary statistics CSV
-"""
-
-import argparse
-import sys
-from pathlib import Path
-
-import matplotlib as mpl
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-from matplotlib.patches import Patch as LegendHandle
-
-SCRIPT_DIR = Path(__file__).resolve().parent
-sys.path.insert(0, str(SCRIPT_DIR))
-
-from pub_utils import (
-    CELL_TYPE_COLORS,
-    load_importance_results,
-    savefig,
-    setup_publication_style,
-)
-
-REACTOME_PREFIX = "module2_reactome"
-
-T_CELL_SUBTYPES = [
-    "CD4+ T Helper2",
-    "CD4+/CD25 T Reg",
-    "CD4+/CD45RA+/CD25- Naive T",
-    "CD4+/CD45RO+ Memory",
-    "CD8+ Cytotoxic T",
-    "CD8+/CD45RA+ Naive Cytotoxic",
-]
-
-T_CELL_SHORT_NAMES = {
-    "CD4+ T Helper2": "CD4+ Th2",
-    "CD4+/CD25 T Reg": "CD4+ Treg",
-    "CD4+/CD45RA+/CD25- Naive T": "CD4+ Naive",
-    "CD4+/CD45RO+ Memory": "CD4+ Memory",
-    "CD8+ Cytotoxic T": "CD8+ Cyto",
-    "CD8+/CD45RA+ Naive Cytotoxic": "CD8+ Naive",
-}
-
-
-def clean_reactome_name(module_name: str) -> str:
-    """Extract readable pathway name from full module name."""
-    parts = module_name.split("__", 1)
-    if len(parts) < 2:
-        return module_name
-    pathway = parts[1]
-    pathway = pathway.rsplit(" R-HSA-", 1)[0]
-    return pathway
-
-
-def extract_reactome_importance(importance_csv: pd.DataFrame):
-    """Extract Reactome-only importance sub-matrix for T cell subtypes."""
-    reactome_cols = [c for c in importance_csv.columns if c.startswith(REACTOME_PREFIX)]
-    t_cell_rows = [ct for ct in T_CELL_SUBTYPES if ct in importance_csv.index]
-
-    sub = importance_csv.loc[t_cell_rows, reactome_cols].copy()
-    sub.columns = [clean_reactome_name(c) for c in sub.columns]
-    return sub
-
-
-# ─── Figure 1: Top-N Reactome per T cell subtype ─────────────────────────────
-
-def plot_top_reactome_per_subtype(sub_df, output_dir, top_n=15):
-    """Horizontal bar chart: top-N Reactome pathways per T cell subtype."""
-    setup_publication_style()
-    n_types = len(sub_df.index)
-    fig, axes = plt.subplots(n_types, 1, figsize=(10, 3.2 * n_types))
-    if n_types == 1:
-        axes = [axes]
-
-    for idx, ct in enumerate(sub_df.index):
-        ax = axes[idx]
-        row = sub_df.loc[ct].sort_values(ascending=False)
-        top = row.head(top_n)
-
-        colors = [CELL_TYPE_COLORS.get(ct, "#333")] * len(top)
-        y_pos = np.arange(len(top))
-        ax.barh(y_pos, top.values, color=colors, alpha=0.8, edgecolor="white",
-                linewidth=0.3, height=0.7)
-
-        short = T_CELL_SHORT_NAMES.get(ct, ct)
-        ax.set_yticks(y_pos)
-        ax.set_yticklabels(top.index, fontsize=7)
-        ax.invert_yaxis()
-        ax.set_xlabel("ΔLogit (importance)", fontsize=9)
-        ax.set_title(f"{short}: Top-{top_n} Reactome Pathways", fontsize=10,
-                     color=CELL_TYPE_COLORS.get(ct, "#333"), fontweight="bold")
-        ax.axvline(0, color="gray", linewidth=0.5, linestyle=":")
-
-        for i, v in enumerate(top.values):
-            if v > 0:
-                ax.text(v + 0.001, i, f"{v:.3f}", va="center", fontsize=6.5)
-
-    fig.tight_layout(h_pad=2.0)
-    savefig(fig, output_dir / "fig1_top_reactome_per_tcell")
-    print("  Saved fig1_top_reactome_per_tcell")
-
-
-# ─── Figure 2: Shared vs Unique critical pathways ────────────────────────────
-
-def analyze_shared_unique(sub_df, output_dir, top_n=20):
-    """Identify shared and unique top Reactome pathways across T cell subtypes."""
-    setup_publication_style()
-
-    top_sets = {}
-    for ct in sub_df.index:
-        row = sub_df.loc[ct].sort_values(ascending=False)
-        top_sets[ct] = set(row.head(top_n).index)
-
-    all_top = set()
-    for s in top_sets.values():
-        all_top |= s
-
-    pathway_counts = {}
-    for pw in all_top:
-        cnt = sum(1 for s in top_sets.values() if pw in s)
-        pathway_counts[pw] = cnt
-
-    shared = {pw: cnt for pw, cnt in pathway_counts.items() if cnt >= 3}
-    unique = {pw: cnt for pw, cnt in pathway_counts.items() if cnt == 1}
-
-    shared_sorted = sorted(shared.items(), key=lambda x: -x[1])
-    unique_by_type = {ct: [] for ct in sub_df.index}
-    for pw, cnt in unique.items():
-        for ct, s in top_sets.items():
-            if pw in s:
-                unique_by_type[ct].append(pw)
-
-    rows = []
-    for pw, cnt in pathway_counts.items():
-        present_in = [T_CELL_SHORT_NAMES.get(ct, ct)
-                      for ct, s in top_sets.items() if pw in s]
-        mean_imp = sub_df[pw].mean()
-        max_imp = sub_df[pw].max()
-        max_type = sub_df[pw].idxmax()
-        rows.append({
-            "pathway": pw,
-            "n_subtypes": cnt,
-            "present_in": ", ".join(present_in),
-            "mean_importance": round(mean_imp, 5),
-            "max_importance": round(max_imp, 5),
-            "max_subtype": T_CELL_SHORT_NAMES.get(max_type, max_type),
-            "category": "shared (≥3)" if cnt >= 3 else ("unique" if cnt == 1 else "partial"),
-        })
-    summary_df = pd.DataFrame(rows).sort_values(
-        ["n_subtypes", "mean_importance"], ascending=[False, False]
-    )
-    summary_df.to_csv(output_dir / "tcell_reactome_shared_unique.csv", index=False)
-    print("  Saved tcell_reactome_shared_unique.csv")
-
-    # --- Figure: Upset-style count bar ---
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5), gridspec_kw={"width_ratios": [2, 1]})
-
-    # Left: shared pathways bar
-    if shared_sorted:
-        pw_names = [x[0] for x in shared_sorted]
-        pw_counts = [x[1] for x in shared_sorted]
-        pw_importance = [sub_df[pw].mean() for pw in pw_names]
-
-        display_names = [n[:50] + "..." if len(n) > 50 else n for n in pw_names]
-        y_pos = np.arange(len(pw_names))
-        bars = ax1.barh(y_pos, pw_importance, color="#1B9E77", alpha=0.8,
-                        edgecolor="white", linewidth=0.3, height=0.7)
-        ax1.set_yticks(y_pos)
-        ax1.set_yticklabels(display_names, fontsize=7)
-        ax1.invert_yaxis()
-        ax1.set_xlabel("Mean ΔLogit across T cell subtypes", fontsize=9)
-        ax1.set_title(f"(A) Shared Reactome Pathways (in ≥3 of 6 subtypes, top-{top_n})",
-                      fontsize=10)
-
-        for i, (imp, cnt) in enumerate(zip(pw_importance, pw_counts)):
-            ax1.text(max(imp, 0) + 0.0005, i, f"{cnt}/6", va="center", fontsize=7,
-                     color="#555", fontweight="bold")
-
-    # Right: unique counts per subtype
-    unique_counts = [len(unique_by_type[ct]) for ct in sub_df.index]
-    short_names = [T_CELL_SHORT_NAMES.get(ct, ct) for ct in sub_df.index]
-    colors = [CELL_TYPE_COLORS.get(ct, "#333") for ct in sub_df.index]
-    y_pos2 = np.arange(len(sub_df.index))
-    ax2.barh(y_pos2, unique_counts, color=colors, alpha=0.8, edgecolor="white",
-             linewidth=0.3, height=0.6)
-    ax2.set_yticks(y_pos2)
-    ax2.set_yticklabels(short_names, fontsize=8)
-    ax2.invert_yaxis()
-    ax2.set_xlabel(f"# Unique pathways (in top-{top_n})", fontsize=9)
-    ax2.set_title("(B) Subtype-Specific Pathways", fontsize=10)
-
-    for i, v in enumerate(unique_counts):
-        if v > 0:
-            ax2.text(v + 0.2, i, str(v), va="center", fontsize=8, fontweight="bold")
-
-    fig.tight_layout(w_pad=3.0)
-    savefig(fig, output_dir / "fig2_shared_unique_reactome")
-    print("  Saved fig2_shared_unique_reactome")
-
-    return summary_df
-
-
-# ─── Figure 3: Heatmap of top Reactome pathways across T cell subtypes ───────
-
-def plot_tcell_reactome_heatmap(sub_df, output_dir, top_n=30):
-    """Heatmap: importance of top Reactome pathways across all T cell subtypes."""
-    setup_publication_style()
-
-    mean_importance = sub_df.mean(axis=0).sort_values(ascending=False)
-    top_pathways = mean_importance.head(top_n).index.tolist()
-
-    heatmap_data = sub_df[top_pathways].copy()
-    short_labels = [T_CELL_SHORT_NAMES.get(ct, ct) for ct in heatmap_data.index]
-
-    display_pathways = [n[:55] + "..." if len(n) > 55 else n for n in top_pathways]
-
-    fig, ax = plt.subplots(figsize=(10, 8))
-    vmax = heatmap_data.values.max()
-    vmin = min(0, heatmap_data.values.min())
-
-    im = ax.imshow(heatmap_data.values.T, cmap="YlOrRd", aspect="auto",
-                   vmin=vmin, vmax=vmax)
-
-    ax.set_xticks(range(len(short_labels)))
-    ax.set_xticklabels(short_labels, fontsize=8, rotation=45, ha="right")
-    ax.set_yticks(range(len(display_pathways)))
-    ax.set_yticklabels(display_pathways, fontsize=7)
-
-    for i in range(len(top_pathways)):
-        for j in range(len(short_labels)):
-            val = heatmap_data.values[j, i]
-            color = "white" if val > vmax * 0.65 else "black"
-            ax.text(j, i, f"{val:.3f}", ha="center", va="center",
-                    fontsize=5.5, color=color)
-
-    cb = fig.colorbar(im, ax=ax, fraction=0.02, pad=0.02)
-    cb.set_label("ΔLogit (importance)", fontsize=9)
-    ax.set_title(f"Top-{top_n} Reactome Pathways: T Cell Subtype Importance",
-                 fontsize=11, pad=10)
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig3_tcell_reactome_heatmap")
-    print("  Saved fig3_tcell_reactome_heatmap")
-
-
-# ─── Figure 4: CD4 vs CD8 comparison ─────────────────────────────────────────
-
-def plot_cd4_vs_cd8_comparison(sub_df, output_dir, top_n=20):
-    """Compare most important Reactome pathways between CD4+ and CD8+ T cells."""
-    setup_publication_style()
-
-    cd4_types = [ct for ct in sub_df.index if "CD4" in ct]
-    cd8_types = [ct for ct in sub_df.index if "CD8" in ct]
-
-    cd4_mean = sub_df.loc[cd4_types].mean(axis=0)
-    cd8_mean = sub_df.loc[cd8_types].mean(axis=0)
-
-    cd4_top = set(cd4_mean.sort_values(ascending=False).head(top_n).index)
-    cd8_top = set(cd8_mean.sort_values(ascending=False).head(top_n).index)
-
-    shared = cd4_top & cd8_top
-    cd4_only = cd4_top - cd8_top
-    cd8_only = cd8_top - cd4_top
-
-    all_pathways = sorted(cd4_top | cd8_top,
-                          key=lambda p: -(cd4_mean[p] + cd8_mean[p]))
-
-    fig, ax = plt.subplots(figsize=(12, max(6, len(all_pathways) * 0.28)))
-
-    display_names = [n[:55] + "..." if len(n) > 55 else n for n in all_pathways]
-    y_pos = np.arange(len(all_pathways))
-    width = 0.35
-
-    bars_cd4 = ax.barh(y_pos - width / 2, [cd4_mean[p] for p in all_pathways],
-                       height=width, color="#984EA3", alpha=0.8, label="CD4+ (mean)")
-    bars_cd8 = ax.barh(y_pos + width / 2, [cd8_mean[p] for p in all_pathways],
-                       height=width, color="#FC8D62", alpha=0.8, label="CD8+ (mean)")
-
-    for i, pw in enumerate(all_pathways):
-        if pw in shared:
-            ax.plot(-0.001, i, "s", color="#333", markersize=4)
-        elif pw in cd4_only:
-            ax.plot(-0.001, i, "s", color="#984EA3", markersize=4)
-        else:
-            ax.plot(-0.001, i, "s", color="#FC8D62", markersize=4)
-
-    ax.set_yticks(y_pos)
-    ax.set_yticklabels(display_names, fontsize=7)
-    ax.invert_yaxis()
-    ax.set_xlabel("Mean ΔLogit (importance)", fontsize=9)
-    ax.set_title(f"CD4+ vs CD8+ T Cells: Top-{top_n} Reactome Pathways", fontsize=11)
-    ax.axvline(0, color="gray", linewidth=0.5, linestyle=":")
-
-    legend_elements = [
-        LegendHandle(facecolor="#984EA3", alpha=0.8, label=f"CD4+ subtypes (n={len(cd4_types)})"),
-        LegendHandle(facecolor="#FC8D62", alpha=0.8, label=f"CD8+ subtypes (n={len(cd8_types)})"),
-        LegendHandle(facecolor="#333", alpha=0.4, label=f"Shared (n={len(shared)})"),
-    ]
-    ax.legend(handles=legend_elements, loc="lower right", fontsize=8)
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig4_cd4_vs_cd8_reactome")
-    print("  Saved fig4_cd4_vs_cd8_reactome")
-
-    comparison_rows = []
-    for pw in all_pathways:
-        comparison_rows.append({
-            "pathway": pw,
-            "cd4_mean": round(cd4_mean[pw], 5),
-            "cd8_mean": round(cd8_mean[pw], 5),
-            "diff_cd4_minus_cd8": round(cd4_mean[pw] - cd8_mean[pw], 5),
-            "category": "shared" if pw in shared else ("CD4-specific" if pw in cd4_only else "CD8-specific"),
-        })
-    comp_df = pd.DataFrame(comparison_rows)
-    comp_df.to_csv(output_dir / "cd4_vs_cd8_reactome.csv", index=False)
-    print("  Saved cd4_vs_cd8_reactome.csv")
-
-    return len(shared), len(cd4_only), len(cd8_only)
-
-
-# ─── Figure 5: Naive vs Effector/Memory ──────────────────────────────────────
-
-def plot_naive_vs_effector(sub_df, output_dir, top_n=20):
-    """Compare naive T cells vs effector/memory T cells."""
-    setup_publication_style()
-
-    naive_types = [ct for ct in sub_df.index
-                   if "Naive" in ct]
-    effector_types = [ct for ct in sub_df.index
-                      if ct not in naive_types]
-
-    naive_mean = sub_df.loc[naive_types].mean(axis=0)
-    effector_mean = sub_df.loc[effector_types].mean(axis=0)
-
-    diff = effector_mean - naive_mean
-    diff_sorted = diff.sort_values()
-
-    bottom_n = diff_sorted.head(top_n)
-    top_n_vals = diff_sorted.tail(top_n)
-    selected = pd.concat([bottom_n, top_n_vals])
-
-    fig, ax = plt.subplots(figsize=(10, max(6, len(selected) * 0.25)))
-
-    display_names = [n[:55] + "..." if len(n) > 55 else n for n in selected.index]
-    y_pos = np.arange(len(selected))
-    colors = ["#2166AC" if v < 0 else "#B2182B" for v in selected.values]
-
-    ax.barh(y_pos, selected.values, color=colors, alpha=0.8,
-            edgecolor="white", linewidth=0.3, height=0.7)
-    ax.set_yticks(y_pos)
-    ax.set_yticklabels(display_names, fontsize=7)
-    ax.axvline(0, color="black", linewidth=0.8)
-    ax.set_xlabel("ΔImportance (Effector/Memory − Naive)", fontsize=9)
-    ax.set_title(f"Reactome Pathways: Effector/Memory vs Naive T Cells (top {top_n} each direction)",
-                 fontsize=10)
-
-    legend_elements = [
-        LegendHandle(facecolor="#B2182B", alpha=0.8, label="More important for Effector/Memory"),
-        LegendHandle(facecolor="#2166AC", alpha=0.8, label="More important for Naive"),
-    ]
-    ax.legend(handles=legend_elements, loc="lower right", fontsize=8)
-
-    fig.tight_layout()
-    savefig(fig, output_dir / "fig5_naive_vs_effector_reactome")
-    print("  Saved fig5_naive_vs_effector_reactome")
-
-    rows = []
-    for pw in diff_sorted.index:
-        rows.append({
-            "pathway": pw,
-            "naive_mean": round(naive_mean[pw], 5),
-            "effector_memory_mean": round(effector_mean[pw], 5),
-            "diff_effector_minus_naive": round(diff[pw], 5),
-        })
-    diff_df = pd.DataFrame(rows)
-    diff_df.to_csv(output_dir / "naive_vs_effector_reactome.csv", index=False)
-    print("  Saved naive_vs_effector_reactome.csv")
-
-
-def parse_args():
-    p = argparse.ArgumentParser(description="T cell Reactome pathway importance analysis")
-    p.add_argument("--results_dir", required=True,
-                   help="Importance results directory (with importance_matrix.csv)")
-    p.add_argument("--output_dir", required=True)
-    p.add_argument("--top_n", type=int, default=15,
-                   help="Number of top pathways per subtype (default: 15)")
-    p.add_argument("--top_n_shared", type=int, default=20,
-                   help="Top-N for shared/unique analysis (default: 20)")
-    p.add_argument("--top_n_heatmap", type=int, default=30,
-                   help="Top-N for heatmap (default: 30)")
-    return p.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    print(f"Loading results from {args.results_dir}...")
-    data = load_importance_results(Path(args.results_dir))
-    importance_csv = data["importance_csv"]
-
-    print(f"Cell types: {list(importance_csv.index)}")
-    t_cell_present = [ct for ct in T_CELL_SUBTYPES if ct in importance_csv.index]
-    print(f"T cell subtypes found: {t_cell_present}")
-
-    sub_df = extract_reactome_importance(importance_csv)
-    print(f"Reactome sub-matrix: {sub_df.shape[0]} T cell types × {sub_df.shape[1]} pathways")
-    print()
-
-    print("[1/5] Top Reactome pathways per T cell subtype...")
-    plot_top_reactome_per_subtype(sub_df, output_dir, top_n=args.top_n)
-    top_per_type_rows = []
-    for ct in sub_df.index:
-        row = sub_df.loc[ct].sort_values(ascending=False)
-        for rank, (pw, val) in enumerate(row.head(args.top_n).items(), 1):
-            top_per_type_rows.append({
-                "cell_type": T_CELL_SHORT_NAMES.get(ct, ct),
-                "rank": rank,
-                "pathway": pw,
-                "importance": round(val, 5),
-            })
-    top_df = pd.DataFrame(top_per_type_rows)
-    top_df.to_csv(output_dir / "top_reactome_per_tcell.csv", index=False)
-    print("  Saved top_reactome_per_tcell.csv")
-    print()
-
-    print("[2/5] Shared vs unique pathway analysis...")
-    summary_df = analyze_shared_unique(sub_df, output_dir, top_n=args.top_n_shared)
-    n_shared = len(summary_df[summary_df["category"] == "shared (≥3)"])
-    n_unique = len(summary_df[summary_df["category"] == "unique"])
-    print(f"  Shared (≥3 subtypes): {n_shared} pathways")
-    print(f"  Unique (1 subtype): {n_unique} pathways")
-    print()
-
-    print("[3/5] T cell Reactome importance heatmap...")
-    plot_tcell_reactome_heatmap(sub_df, output_dir, top_n=args.top_n_heatmap)
-    print()
-
-    print("[4/5] CD4+ vs CD8+ comparison...")
-    n_sh, n_cd4, n_cd8 = plot_cd4_vs_cd8_comparison(sub_df, output_dir, top_n=args.top_n_shared)
-    print(f"  Shared: {n_sh}, CD4-specific: {n_cd4}, CD8-specific: {n_cd8}")
-    print()
-
-    print("[5/5] Naive vs Effector/Memory comparison...")
-    plot_naive_vs_effector(sub_df, output_dir, top_n=args.top_n)
-    print()
-
-    print("All analyses complete.")
-    print(f"Output directory: {output_dir}")
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_text_alignment/DESIGN.md

@@ -1,235 +0,0 @@
-# Experiment 2: Module Embedding to Pathway Text Alignment
-
-## Overview
-
-This experiment tests whether the structure of the BioMod module embedding space
-significantly aligns with the semantic space of natural-language pathway
-descriptions. Gene-level models cannot reproduce this experiment because single
-gene names do not carry phrase-level semantics.
-
-## Scripts
-
-| Script | Function |
-| --- | --- |
-| `biomod_text_alignment.py` | Main script: RSA/Mantel test, CKA, and leave-one-out Ridge retrieval |
-| `biomod_text_alignment_extended.py` | Extended analysis: Reactome broad-category RSA, Reactome description-based RSA, and Hallmark deep dive |
-| `biomod_text_alignment_publication.py` | Publication-level analysis: six main figures, bootstrap CIs, multi-model comparison, cross-dataset validation, and deep analysis |
-| `fetch_reactome_descriptions.py` | Fetch `summaryText` for 491 pathways from the Reactome API |
-
-## Commands
-
-Run from the BioMod repository root. Generate each `module_embeddings.npy` with
-`BioMod/interpretability/biomod_geometry.py` before using it here. The historical
-`pretrain_subset` output requires a separately supplied pretraining dataset;
-the `pancreas_scib` example uses the bundled pancreas h5ad file. Existing result
-tables below describe historical runs; their raw output directories are not
-included in this repository.
-
-### 1. Download the Sentence Encoder
-
-```bash
-# Run from the BioMod repository root.
-# Activate the environment configured for this project.
-
-# The sentence-transformers loader downloads this model on first use.
-# Network access or an existing local model cache is required.
-
-pip install sentence-transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
-```
-
-### 2. Main Experiment: RSA + CKA + LOO Retrieval
-
-```bash
-# pretrain_subset is recommended because averaging over 50k cells is more stable.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --n_permutations 1000 \
-    --output_dir outputs/interpretability/text_alignment/pretrain_subset_v2
-
-# pancreas_scib validates cross-dataset consistency.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pancreas_scib/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --n_permutations 1000 \
-    --output_dir outputs/interpretability/text_alignment/pancreas_scib_v2
-
-# Quick mode: skip LOO retrieval and run only RSA + CKA.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --skip_loo \
-    --output_dir outputs/interpretability/text_alignment/pretrain_subset_quick
-```
-
-### 3. Fetch Official Reactome Descriptions (Completed, 486/491)
-
-```bash
-# Re-fetch if needed. Requires network access to reactome.org and takes about 2 minutes.
-python BioMod/interpretability/biomod_text_alignment/fetch_reactome_descriptions.py \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --output reactome_descriptions.json --delay 0.05
-```
-
-The command writes refreshed descriptions to `reactome_descriptions.json` in the
-repository root. The analysis examples below use the bundled
-`BioMod/interpretability/biomod_text_alignment/reactome_descriptions.json` instead;
-set `--descriptions_file reactome_descriptions.json` to use the refreshed output.
-
-### 4. Extended Analysis: Reactome Categories, Description-Based RSA, Hallmark Deep Dive
-
-```bash
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_extended.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --descriptions_file BioMod/interpretability/biomod_text_alignment/reactome_descriptions.json \
-    --n_permutations 1000 \
-    --output_dir outputs/interpretability/text_alignment/pretrain_subset_v2_extended
-```
-
-### 5. Publication-Level Analysis
-
-```bash
-# Basic run: one dataset and one model.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_publication.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --descriptions_file BioMod/interpretability/biomod_text_alignment/reactome_descriptions.json \
-    --n_permutations 1000 \
-    --n_bootstrap 2000 \
-    --output_dir outputs/interpretability/text_alignment/publication
-
-# Full run: cross-dataset validation and multi-model comparison.
-python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_publication.py \
-    --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-    --biomod_embedding_path2 outputs/interpretability/geometry/pancreas_scib/module_embeddings.npy \
-    --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-    --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-    --sentence_models_compare "sentence-transformers/all-mpnet-base-v2,sentence-transformers/all-MiniLM-L6-v2" \
-    --descriptions_file BioMod/interpretability/biomod_text_alignment/reactome_descriptions.json \
-    --n_permutations 1000 \
-    --n_bootstrap 2000 \
-    --output_dir outputs/interpretability/text_alignment/publication_full
-```
-
-**Arguments:**
-
-| Argument | Description |
-| --- | --- |
-| `--biomod_embedding_path` | `module_embeddings.npy` for the main dataset, from Experiment 1 |
-| `--biomod_embedding_path2` | Optional second dataset for cross-dataset validation |
-| `--sentence_models_compare` | Comma-separated model names for comparison, downloaded automatically from HuggingFace |
-| `--n_permutations` | Mantel-test permutations; 1000 is standard and 10000 gives more precise p-values |
-| `--n_bootstrap` | Bootstrap resampling iterations; 2000 is standard for 95% CIs |
-| `--device` | Encoding device, `cpu` or `cuda`; default is CPU |
-
-## Dependencies
-
-```bash
-pip install sentence-transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
-```
-
-## Output Files
-
-### Main Experiment Outputs (`biomod_text_alignment.py`)
-
-| File | Content |
-| --- | --- |
-| `pathway_descriptions.json` | Text descriptions for 541 modules |
-| `text_embeddings.npy` | Sentence-encoder embeddings with shape `[541, 768]` |
-| `module_embeddings_subset.npy` | Corresponding model module embeddings with shape `[541, 1024]` |
-| `alignment_metrics.json` | RSA, CKA, retrieval recall@k, and p-values |
-| `rsa_scatter.png` | Paired-distance scatter plot: module space vs text space |
-| `rsa_null_distribution.png` | Mantel test: observed value vs null distribution |
-| `retrieval_ranks.png` | Histogram of retrieval ranks |
-| `retrieval_examples.csv` | Top-5 retrieval results for representative pathways |
-| `run.log` | Full log |
-
-### Extended Analysis Outputs (`biomod_text_alignment_extended.py`)
-
-| File | Content |
-| --- | --- |
-| `extended_metrics.json` | Summary metrics for the three analyses |
-| `reactome_category_heatmaps.png` | Broad Reactome category distance-matrix comparison, module vs text |
-| `hallmarks_deep_dive.png` | Hallmark similarity matrices and per-module alignment-score bar chart |
-| `hallmarks_per_module_alignment.csv` | Per-module alignment ranking for 50 Hallmarks |
-| `run_extended.log` | Extended-analysis log |
-
-### Publication-Level Outputs (`biomod_text_alignment_publication.py`)
-
-| File | Content |
-| --- | --- |
-| `fig1_rsa_overview.{pdf,png}` | 2D hexbin density plus Mantel null distribution with bootstrap CI |
-| `fig2_hallmarks_heatmaps.{pdf,png}` | Three 50 x 50 similarity matrices: module, text, and difference; hierarchical clustering order and adaptive color scale |
-| `fig3_per_hallmark_alignment.{pdf,png}` | Per-Hallmark alignment-score lollipop plot, colored by MSigDB process category |
-| `fig4_category_alignment.{pdf,png}` | Box + strip plot grouped by biological process category with Kruskal-Wallis test |
-| `fig5_granularity_effect.{pdf,png}` | Bar plot of RSA across granularity levels: fine-grained, category, superclass, Hallmark |
-| `fig6_model_comparison.{pdf,png}` | Optional RSA comparison across models and datasets |
-| `fig_alignment_vs_size.{pdf,png}` | Alignment score vs module gene count scatter plot |
-| `publication_metrics.json` | Summary of all metrics, including bootstrap 95% CIs |
-| `hallmarks_alignment_table.csv` | Alignment scores and process categories for 50 Hallmarks |
-| `negative_analysis.csv` | Neighbor comparison for the worst- and best-aligned modules |
-| `run_publication.log` | Full log |
-
-## Results (v2, pretrain_subset)
-
-### Main Experiment
-
-| Metric | All (541) | Hallmarks (50) | Reactome (491) |
-| --- | --- | --- | --- |
-| RSA Spearman r | 0.050 | **0.160** | -0.007 |
-| Mantel p-value | 0.012 | 0.030 | 0.586 |
-| z-score | 2.2 sigma | 2.1 sigma | -0.3 sigma |
-| CKA | 0.043 (p=0.000) | - | - |
-| Retrieval median rank | 232 (random=270) | - | - |
-
-### Interpretation
-
-- **Hallmark alignment is significant (r=0.16, p=0.03):** among 50 broad biological processes, distance structure in model space significantly correlates with semantic distance in text space.
-- **Reactome does not align (r=-0.007, p=0.59):** among 491 fine-grained pathways there is no alignment, likely because fine-grained pathway gene-composition similarity, encoded by the model, and pathway-name language similarity, encoded by text, live on different axes.
-- **Global CKA is significant (p=0.000):** the overall representation structure contains an alignment signal.
-
-### Improvement History
-
-**v1, BioLinkBERT + 1-to-1 retrieval:**
-
-- BioLinkBERT was used as a raw BERT model with mean pooling rather than a sentence-level fine-tuned encoder.
-- The main metric was exact 1-to-1 retrieval, which is too strict for a 541-dimensional space.
-- CKA=0.047 was statistically significant but weak; retrieval was close to random.
-
-**v2, S-PubMedBert-MS-MARCO + RSA/Mantel test:**
-
-1. Switched to a sentence encoder: PubMedBERT fine-tuned on MS-MARCO sentence similarity.
-2. Added RSA/Mantel test as the main metric, testing structural rather than exact retrieval alignment.
-3. Per-family RSA revealed that the alignment signal is concentrated in Hallmarks.
-
-### Extended Analysis Directions
-
-1. **Reactome broad-category aggregation:** aggregate the 491 pathways into about 16 top-level categories and retest RSA; broad-category alignment is expected to be significant.
-2. **Reactome descriptions:** use API-fetched `summaryText` as richer text input and test whether alignment improves.
-3. **Hallmark deep dive:** compute per-module alignment scores, kNN overlap, and the best/worst aligned Hallmarks.
-
-## Design Decisions
-
-### Why Use RSA Instead of 1-to-1 Retrieval, 
-
-1-to-1 retrieval asks whether the paired text can be found exactly from a module
-embedding, requiring a very strong cross-modal mapping. RSA asks whether modules
-that are close in model space are also close in text space, which is a more
-natural definition of structural alignment.
-
-### Why Does Reactome Not Align, 
-
-- Reactome pathways overlap heavily because parent and child pathways share many genes.
-- The model encodes gene-composition similarity, while text names reflect hierarchy and naming similarity; the two do not perfectly match.
-- The granularity is too fine. For example, text distance between `Signaling by FGFR1` and `Signaling by FGFR2` does not necessarily reflect biological distance.
-
-### Why Exclude Hotspot/Processed, 
-
-- Hotspot modules (73) are data-driven co-expression modules with only numeric IDs and no pathway semantics.
-- Processed modules (46) are pipeline-derived gene groups such as `SM_2` and have no pathway names.
-- Without text pairs, they cannot participate in this alignment experiment.

BioMod/interpretability/biomod_text_alignment/MANUSCRIPT_SECTION.md

@@ -1,208 +0,0 @@
-# Experiment 2: Module Embedding to Pathway Text Alignment
-
-> This document contains a Results subsection draft, complete figure captions,
-> and a supplementary-material framework for an NC submission. The main text is
-> written in English.
-
----
-
-## Results Subsection Draft
-
-### Module embeddings align with pathway semantic structure
-
-To test whether scLong's learned representations capture biologically meaningful relationships beyond gene-level co-expression, we examined whether the geometric structure of module embeddings aligns with the semantic structure of pathway natural language descriptions, an experiment that is structurally impossible for gene-level foundation models because individual gene names lack phrase-level semantics.
-
-We encoded the names and descriptions of 541 pathway-based modules (50 MSigDB Hallmarks + 491 Reactome pathways) using S-PubMedBert-MS-MARCO, a biomedical sentence encoder fine-tuned for semantic similarity. We then quantified the structural correspondence between the module embedding space and the text embedding space using Representational Similarity Analysis (RSA) with a Mantel permutation test (1,000 permutations).
-
-The module embedding space showed significant structural alignment with the pathway text space (Spearman r = 0.41, p < 0.001, z = 14.8 sigma; Fig. Xa). Bootstrap resampling confirmed the robustness of this estimate (95% CI: [0.32, 0.49]). Critically, this alignment was reproduced on an independent dataset (pancreas scIB; r = 0.40, z = 14.1 sigma, p < 0.001), demonstrating that the observed correspondence is not dataset-specific but reflects a stable property of the learned representation.
-
-**Granularity-dependent alignment.** We observed a striking granularity effect: alignment was absent at the fine-grained Reactome pathway level (491 pathways; r = 0.001, p > 0.5), moderate at the Reactome top-level category level (17 categories; r = 0.27, 95% CI: [0.06, 0.68]), and significant at the Hallmarks level (50 pathways; r = 0.16, p = 0.03). This pattern is biologically expected because fine-grained Reactome pathways share extensive gene overlap, making their gene-composition similarity orthogonal to their name similarity. At coarser granularity, pathway names better reflect the functional distinctions that the model has learned.
-
-**Category-specific alignment patterns.** Per-hallmark alignment analysis revealed that 38 of 50 hallmarks (76%) showed positive alignment scores (Fig. Xc). Grouping by biological process category revealed significant between-category differences (Kruskal-Wallis H = 20.8, p = 0.004; Fig. Xd). Signaling pathways (median r = 0.33) and proliferation-related pathways (median r = 0.32) showed the strongest alignment, while metabolic (median r = -0.05) and developmental (median r = -0.04) pathways showed weak or negative alignment. This hierarchy suggests that the model most faithfully encodes functional relationships for pathways with coherent, non-overlapping gene programs.
-
-**Nearest-neighbor overlap.** As a complementary metric, we computed k-nearest-neighbor overlap (k = 5) between the two spaces. The mean overlap was 0.20, significantly exceeding the permutation null of 0.10 (z = 4.5 sigma, p < 0.001; Extended Data Fig. Xc).
-
-**Informative negative results.** The five least-aligned hallmarks (Apoptosis, r = -0.39; Inflammatory Response, r = -0.27; EMT, r = -0.26; Fatty Acid Metabolism, r = -0.26; Spermatogenesis, r = -0.14) provide biological insight. For example, Apoptosis is embedded near Apical Surface and Bile Acid Metabolism in the model space, reflecting shared gene composition, but near Adipogenesis and Angiogenesis in text space, reflecting name-level semantics. This divergence shows that the model encodes functional gene-composition relationships rather than superficial naming conventions, which is desirable for a biological foundation model.
-
-**Independence from module size.** Alignment scores showed no correlation with the number of genes per module (Spearman r = -0.17, p = 0.23; Extended Data Fig. Xa), ruling out the trivial explanation that larger modules produce more stable embeddings and thus better alignment.
-
----
-
-## Main Figure Caption
-
-### Figure X. Module embedding space aligns with pathway semantic structure.
-
-**(a)** Representational Similarity Analysis (RSA) between module embedding distances and pathway text embedding distances for 541 pathway-based modules. Left: 2D hexbin density plot showing the correlation between pairwise cosine distances in module space (x-axis) and text space (y-axis). Spearman r = 0.41, p < 0.001 (Mantel test, 1,000 permutations). Right: null distribution from permutation test (gray histogram) with observed correlation (red dashed line) and 95% bootstrap confidence interval (red shaded region). z = 14.8 sigma.
-
-**(b)** Rank-transformed cosine similarity matrices for the 50 MSigDB Hallmark modules, ordered by hierarchical clustering (Ward's method on module embeddings). Left: module embedding similarity. Middle: text embedding similarity. Right: difference (module - text) in rank space. Color bar above indicates MSigDB biological process category. Rank transformation converts the narrow raw cosine similarity range ([0.988, 1.000] for module; [0.76, 1.00] for text) to a uniform [0, 1] scale, revealing block-diagonal structure corresponding to functionally related pathway clusters.
-
-**(c)** Per-hallmark alignment scores (Spearman correlation between each hallmark's similarity profile in module space vs. text space), displayed as a lollipop chart sorted by score. Colors indicate biological process category. 38/50 hallmarks (76%) show positive alignment. Mean alignment = 0.12 (95% CI: [0.07, 0.18]).
-
-**(d)** Alignment scores grouped by MSigDB biological process category. Box-and-strip plot with individual hallmark scores overlaid. Kruskal-Wallis H = 20.8, p = 0.004. Signaling and proliferation categories show significantly stronger alignment than metabolic and developmental categories.
-
----
-
-## Extended Data Figure Caption
-
-### Extended Data Figure X. Supplementary module-text alignment analyses.
-
-**(a)** Alignment score vs. number of genes per module. Spearman r = -0.17, p = 0.23, indicating no dependence on module size.
-
-**(b)** RSA as a function of pathway granularity. Bar chart showing Spearman r at four levels: fine-grained Reactome (n = 491), Reactome top-level categories (n = 17), MSigDB Hallmarks (n = 50), and Reactome super-categories (n = 6). Error bars: 95% bootstrap CI. Alignment peaks at intermediate granularity where pathway names best reflect functional distinctions.
-
-**(c)** Negative result analysis: top-5 and bottom-5 aligned hallmarks with their 5 nearest neighbors in module space vs. text space. Divergent neighbors show that the model encodes gene-composition similarity rather than name-level semantics.
-
----
-
-## Key Statistics Summary
-
-| Metric | Value | 95% CI | p-value | Note |
-| --- | --- | --- | --- | --- |
-| RSA (all 541, primary) | 0.41 | [0.32, 0.49] | < 0.001 | z = 14.8 sigma |
-| RSA (all 541, secondary/pancreas) | 0.40 | [0.32, 0.48] | < 0.001 | z = 14.1 sigma |
-| RSA (Hallmarks 50) | 0.16 | [0.05, 0.39] | 0.026 | - |
-| RSA (Reactome fine 491) | 0.001 | [-0.04, 0.06] | > 0.5 | Expected null |
-| RSA (Reactome category 17) | 0.27 | [0.06, 0.68] | - | Wide CI due to n=17 |
-| kNN overlap (k=5) | 0.20 | [0.16, 0.24] | < 0.001 | null = 0.10, z = 4.5 sigma |
-| Per-hallmark positive rate | 38/50 (76%) | - | - | - |
-| Mean per-hallmark alignment | 0.12 | [0.07, 0.18] | - | - |
-| Category effect (Kruskal-Wallis) | H = 20.8 | - | 0.004 | 8 categories |
-| Alignment vs module size | r = -0.17 | - | 0.23 | No correlation |
-
-### Top-5 Best-Aligned Hallmarks
-
-| Hallmark | Alignment r | Category |
-| --- | --- | --- |
-| G2M Checkpoint | 0.47 | proliferation |
-| KRAS Signaling Up | 0.46 | signaling |
-| KRAS Signaling Dn | 0.42 | signaling |
-| MYC Targets V2 | 0.42 | proliferation |
-| mTORC1 Signaling | 0.35 | signaling |
-
-### Bottom-5 Worst-Aligned Hallmarks
-
-| Hallmark | Alignment r | Module NN | Text NN |
-| --- | --- | --- | --- |
-| Apoptosis | -0.39 | Apical Surface, Bile Acid, Cholesterol, Coagulation | Adipogenesis, Angiogenesis, Spermatogenesis |
-| Inflammatory Response | -0.27 | Hypoxia, KRAS Up, KRAS Dn, MYC V1 | Androgen, IFN gamma, IFN alpha, UPR, TNF alpha |
-| EMT | -0.26 | G2M, Fatty Acid, Estrogen Late, E2F | Apical Junction, Inflammatory, Estrogen Early |
-| Fatty Acid Metabolism | -0.26 | Estrogen Late, G2M, Hypoxia, Hedgehog | Bile Acid, Heme, Glycolysis, Xenobiotic |
-| Spermatogenesis | -0.14 | UV Dn, Myogenesis, Notch, Xenobiotic | Adipogenesis, Apoptosis, Peroxisome |
-
----
-
-## Methods Subsection Draft
-
-### Module embedding-pathway text alignment (RSA)
-
-To quantify the structural correspondence between scLong's module embedding space and pathway semantic space, we performed Representational Similarity Analysis (RSA) with a Mantel permutation test.
-
-**Text encoding.** For each of the 541 pathway-based modules (50 MSigDB Hallmarks and 491 Reactome pathways), we constructed a text description consisting of the pathway name and, where available, the Reactome summaryText (obtained via the Reactome Content Service API; 486/491 pathways had descriptions). Texts were encoded using S-PubMedBert-MS-MARCO (PubMedBERT fine-tuned on MS-MARCO for sentence-level semantic similarity), producing 768-dimensional text embeddings.
-
-**Module embeddings.** Cell-averaged module embeddings (1,024-dimensional) were obtained from the module geometry analysis by running the pretrained scLong encoder with mask_ratio = 0 on 50,000 cells from the pretraining corpus, then averaging across cells.
-
-**RSA and Mantel test.** We computed pairwise cosine distance matrices for both the module embeddings (541 x 541) and text embeddings (541 x 541), then measured their Spearman rank correlation. Statistical significance was assessed via a Mantel permutation test: we shuffled the row/column assignment of one distance matrix 1,000 times to construct a null distribution of correlations under the hypothesis of no structural correspondence. The p-value was computed as the fraction of permuted correlations exceeding the observed value. Effect size was reported as z = (observed - null_mean) / null_std.
-
-**Bootstrap confidence intervals.** We computed 95% confidence intervals for RSA by resampling modules with replacement (2,000 iterations) and recomputing the Spearman correlation on each bootstrap sample.
-
-**Per-hallmark alignment scores.** For each of the 50 Hallmark modules, we computed the Spearman correlation between its similarity profile in module space (cosine similarities to all other 49 Hallmarks) and its similarity profile in text space. This yields a per-pathway alignment score ranging from -1, anti-aligned, to +1, perfectly aligned.
-
-**kNN overlap.** For each module, we identified its k = 5 nearest neighbors in both spaces and computed the fraction of shared neighbors. The mean overlap was compared to a permutation null with 1,000 shuffles.
-
-**Cross-dataset validation.** We repeated the full RSA analysis using module embeddings derived from an independent dataset (pancreas scIB, about 15,000 cells) to assess generalizability.
-
-**Granularity analysis.** We computed RSA at four levels of pathway granularity: individual Reactome pathways (n = 491), Reactome top-level categories (n = 17, obtained by mapping each pathway to its root ancestor in the Reactome hierarchy), MSigDB Hallmarks (n = 50), and Reactome super-categories (n = 6, by merging related top-level categories). For category-level analysis, module embeddings and text embeddings were averaged within each category before computing pairwise distances.
-
----
-
-## Manuscript Narrative Logic
-
-### Core Claim
-
-The scLong module embedding space spontaneously aligns with pathway semantic
-space. This is a unique advantage of module-level design:
-
-1. **Gene-level models cannot reproduce it** because single gene names lack phrase-level semantics and cannot support RSA with pathway descriptions.
-2. **The alignment is learned**; combined with the init/shuffle controls from Experiment 1, it supports a pretraining-derived signal.
-3. **Alignment appears at the right granularity**, matching biological expectation. Too fine a level is dominated by gene overlap; too coarse a level has too few samples.
-4. **Negative results are informative**, because non-aligned pathways show that the model encodes function rather than names.
-5. **Cross-dataset consistency** rules out dataset-specific artifacts.
-
-### Connection to Experiment 1
-
-Experiment 1 shows that the module space self-organizes by pathway family, which
-is geometric structure. Experiment 2 further shows that this structure aligns
-with a human-readable semantic space, providing semantic interpretability. The
-two experiments jointly support intrinsic biological interpretability in scLong
-module representations.
-
-### Differentiation from Baselines
-
-> Gene-level foundation models (scGPT, Geneformer, scFoundation) embed individual genes,
-> whose names are arbitrary identifiers such as `BRCA1` and `TP53` and lack compositional semantics.
-> Consequently, the RSA framework applied here, which requires phrase-level text descriptions
-> to construct a meaningful semantic distance matrix, is structurally inapplicable to gene-level
-> models. This represents a qualitative interpretability advantage of the module-level design.
-
-### How to Use Negative Results in the Main Text
-
-Negative results, such as metabolic/development pathways not aligning, are not a
-weakness but evidence:
-
-- The model encodes **functional similarity in gene composition**, not linguistic name similarity.
-- This is exactly what a strong biological model should do.
-- If the model only learned name similarity, the result would be trivial.
-
----
-
-## Figure File Mapping
-
-### Main Figure
-
-| Panel | File | Content |
-| --- | --- | --- |
-| (a) | `fig1_rsa_overview.pdf` | Hexbin density plot and Mantel null distribution |
-| (b) | `fig2_hallmarks_heatmaps.pdf` | Rank-transformed three-panel heatmap and dendrogram |
-| (c) | `fig3_per_hallmark_alignment.pdf` | Lollipop chart colored by category |
-| (d) | `fig4_category_alignment.pdf` | Box + strip plot with Kruskal-Wallis test |
-
-### Extended Data
-
-| Panel | File | Content |
-| --- | --- | --- |
-| (a) | `fig_alignment_vs_size.pdf` | Alignment vs module gene count |
-| (b) | `fig5_granularity_effect.pdf` | RSA across four granularity levels |
-| (c) | `negative_analysis.csv` | Neighbor comparison table |
-| Table | `hallmarks_alignment_table.csv` | Full ranking for 50 Hallmarks |
-| Metrics | `publication_metrics.json` | Machine-readable summary |
-
----
-
-## Category Median Details for Main Text
-
-| Category | n | Median r | Interpretation |
-| --- | --- | --- | --- |
-| Signaling | 11 | 0.333 | Pathway names and functions are highly consistent |
-| Proliferation | 6 | 0.323 | Proliferation-related modules have clear semantics |
-| DNA Damage | 3 | 0.215 | Good |
-| Pathway | 5 | 0.182 | Moderate |
-| Cellular Component | 3 | 0.170 | Moderate |
-| Immune | 9 | 0.121 | Moderately low |
-| Development | 5 | -0.041 | Not aligned |
-| Metabolic | 8 | -0.052 | Not aligned |
-
----
-
-## kNN Overlap by Category
-
-| Category | Mean Overlap | Interpretation |
-| --- | --- | --- |
-| Cellular Component | 0.40 | Highest; cellular-component modules share neighbors across both spaces |
-| Immune | 0.29 | High; immune pathways cluster stably |
-| Metabolic | 0.20 | Moderate |
-| Proliferation | 0.17 | Moderate |
-| Signaling | 0.16 | Moderate |
-| Development | 0.16 | Moderate |
-| DNA Damage | 0.13 | Low |
-| Pathway | 0.12 | Low |
-| **Null (random)** | **0.10** | Baseline |

BioMod/interpretability/biomod_text_alignment/biomod_text_alignment.py

@@ -1,665 +0,0 @@
-"""Experiment 2 — Module embedding <-> pathway text representation alignment.
-
-Question
---------
-Do scLong's module embeddings align with the natural-language semantics of
-their corresponding pathway names/descriptions,  If yes, this demonstrates
-that the model has learned biologically meaningful representations at the
-module level — something gene-level baselines cannot replicate because
-individual gene names lack phrase-level semantics.
-
-Pipeline
---------
-1. Parse pathway names from merged_biomods.json keys for Hallmarks (50) and
-   Reactome (491) modules. Hotspot/Processed are excluded (no semantic names).
-2. Encode pathway names with a biomedical sentence encoder.
-3. Load cell-averaged module embeddings from experiment 1 output.
-4. Measure alignment:
-   - RSA (Representational Similarity Analysis) with Mantel test
-   - CKA (Centered Kernel Alignment) with permutation test
-   - Ridge regression linear mapping (LOO-CV) for retrieval
-5. Visualize: RSA scatter, similarity distribution, retrieval examples.
-
-Data requirement
-----------------
-Only needs module_embeddings.npy from experiment 1 (any dataset). No raw
-expression data needed — this experiment operates entirely in embedding space.
-
-Sentence encoder
-----------------
-Default: pritamdeka/S-PubMedBert-MS-MARCO (a PubMedBERT fine-tuned for
-sentence similarity via MS-MARCO). Unlike raw BioLinkBERT (which is just a
-masked LM without sentence-level training), this model produces embeddings
-optimized for semantic similarity between sentences.
-
-Alternative: sentence-transformers/all-mpnet-base-v2 (general-domain, strong).
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-from scipy.spatial.distance import cosine as cosine_dist
-from scipy.spatial.distance import pdist, squareform
-from scipy.stats import spearmanr, pearsonr
-from sklearn.linear_model import RidgeCV
-from sklearn.metrics.pairwise import cosine_similarity
-from sklearn.preprocessing import StandardScaler
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-HERE = Path(__file__).resolve().parent
-INTERP_DIR = HERE.parent
-BIOMOD_SOURCE = INTERP_DIR.parent
-REPO_ROOT = BIOMOD_SOURCE.parent
-for p in (str(INTERP_DIR), str(BIOMOD_SOURCE), str(REPO_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from interpretability.utils import (
-    MODULE_FAMILY_MAP,
-    load_biomod_names,
-    parse_biomod_family,
-)
-
-logger = logging.getLogger(__name__)
-
-
-# ============================================================
-# 1. Pathway name parsing
-# ============================================================
-
-def parse_pathway_text(module_name: str) -> Optional[str]:
-    """Extract a human-readable pathway name from a module key."""
-    family = parse_biomod_family(module_name)
-    if family == "Hallmarks":
-        raw = module_name.split("__", 1)[1]
-        raw = raw.replace("HALLMARK_", "", 1)
-        return raw.replace("_", " ").title()
-    elif family == "Reactome":
-        raw = module_name.split("__", 1)[1]
-        parts = raw.rsplit(" R-HSA-", 1)
-        return parts[0].strip()
-    return None
-
-
-def build_text_descriptions(
-    module_names: List[str],
-    descriptions_file: Optional[str] = None,
-) -> Tuple[List[int], List[str], List[str]]:
-    """Build text descriptions for modules that have semantic names.
-
-    Returns: (indices, names, texts) — only for Hallmarks + Reactome modules.
-    """
-    ext_descriptions = {}
-    if descriptions_file and Path(descriptions_file).exists():
-        with open(descriptions_file, "r", encoding="utf-8") as f:
-            ext_descriptions = json.load(f)
-
-    indices = []
-    names = []
-    texts = []
-    for i, pname in enumerate(module_names):
-        text = parse_pathway_text(pname)
-        if text is None:
-            continue
-        indices.append(i)
-        names.append(pname)
-        if pname in ext_descriptions:
-            texts.append(f"{text}: {ext_descriptions[pname]}")
-        else:
-            texts.append(text)
-    return indices, names, texts
-
-
-# ============================================================
-# 2. Text encoding
-# ============================================================
-
-def encode_texts(
-    texts: List[str],
-    model_name: str = "pritamdeka/S-PubMedBert-MS-MARCO",
-    batch_size: int = 32,
-    device: str = "cpu",
-) -> np.ndarray:
-    """Encode texts using a sentence-transformers model."""
-    try:
-        from sentence_transformers import SentenceTransformer
-    except ImportError:
-        raise ImportError(
-            "sentence-transformers required. Install: pip install sentence-transformers"
-        )
-
-    logger.info("Loading sentence encoder: %s", model_name)
-    model = SentenceTransformer(model_name, device=device)
-    logger.info("Encoding %d texts...", len(texts))
-    embeddings = model.encode(
-        texts, batch_size=batch_size, show_progress_bar=True,
-        normalize_embeddings=True,
-    )
-    return np.asarray(embeddings, dtype=np.float32)
-
-
-# ============================================================
-# 3. RSA (Representational Similarity Analysis)
-# ============================================================
-
-def compute_rsa(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    metric: str = "cosine",
-) -> Tuple[float, float, float, float]:
-    """Representational Similarity Analysis.
-
-    Computes pairwise distance matrices in both spaces, then correlates them.
-    This tests: "are modules that are close in model space also close in text space, "
-
-    Returns: (spearman_r, spearman_p, pearson_r, pearson_p)
-    """
-    dist_module = pdist(module_emb, metric=metric)
-    dist_text = pdist(text_emb, metric=metric)
-
-    sp_r, sp_p = spearmanr(dist_module, dist_text)
-    pe_r, pe_p = pearsonr(dist_module, dist_text)
-
-    return float(sp_r), float(sp_p), float(pe_r), float(pe_p)
-
-
-def mantel_test(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    metric: str = "cosine",
-    n_perm: int = 1000,
-    seed: int = 42,
-) -> Dict[str, float]:
-    """Mantel test: permutation-based significance for RSA.
-
-    Shuffles rows of one matrix and recomputes the distance-matrix correlation.
-    More robust than parametric p-values for distance correlations.
-    """
-    dist_module = pdist(module_emb, metric=metric)
-    dist_text = pdist(text_emb, metric=metric)
-
-    observed_r, _ = spearmanr(dist_module, dist_text)
-
-    rng = np.random.default_rng(seed)
-    null_rs = np.zeros(n_perm)
-    n = module_emb.shape[0]
-    for k in range(n_perm):
-        perm = rng.permutation(n)
-        dist_text_perm = pdist(text_emb[perm], metric=metric)
-        null_rs[k], _ = spearmanr(dist_module, dist_text_perm)
-
-    p_value = float(np.mean(null_rs >= observed_r))
-    return {
-        "rsa_spearman_observed": float(observed_r),
-        "rsa_null_mean": float(null_rs.mean()),
-        "rsa_null_std": float(null_rs.std()),
-        "rsa_p_value": p_value,
-        "rsa_effect_size_z": float((observed_r - null_rs.mean()) / max(null_rs.std(), 1e-10)),
-    }
-
-
-def compute_rsa_by_family(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    metric: str = "cosine",
-    n_perm: int = 500,
-    seed: int = 42,
-) -> Dict[str, Dict[str, float]]:
-    """RSA computed separately for each module family (Hallmarks vs Reactome)."""
-    results = {}
-    families = [parse_biomod_family(n) for n in names]
-    for fam in ["Hallmarks", "Reactome"]:
-        mask = np.array([f == fam for f in families])
-        if mask.sum() < 10:
-            continue
-        sub_module = module_emb[mask]
-        sub_text = text_emb[mask]
-        mantel = mantel_test(sub_module, sub_text, metric=metric,
-                            n_perm=n_perm, seed=seed)
-        results[fam] = mantel
-    return results
-
-
-# ============================================================
-# 4. CKA
-# ============================================================
-
-def compute_cka(X: np.ndarray, Y: np.ndarray) -> float:
-    """Centered Kernel Alignment (linear kernel)."""
-    n = X.shape[0]
-    H = np.eye(n) - np.ones((n, n)) / n
-    Kx = H @ (X @ X.T) @ H
-    Ky = H @ (Y @ Y.T) @ H
-    hsic = np.sum(Kx * Ky)
-    norm = np.sqrt(np.sum(Kx * Kx) * np.sum(Ky * Ky))
-    if norm < 1e-10:
-        return 0.0
-    return float(hsic / norm)
-
-
-def permutation_test_cka(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    n_perm: int = 1000,
-    seed: int = 42,
-) -> Dict[str, float]:
-    """CKA with permutation null distribution."""
-    observed = compute_cka(module_emb, text_emb)
-    rng = np.random.default_rng(seed)
-    null_values = np.zeros(n_perm)
-    for k in range(n_perm):
-        perm = rng.permutation(text_emb.shape[0])
-        null_values[k] = compute_cka(module_emb, text_emb[perm])
-    p_value = float(np.mean(null_values >= observed))
-    return {
-        "cka_observed": observed,
-        "cka_null_mean": float(null_values.mean()),
-        "cka_null_std": float(null_values.std()),
-        "cka_p_value": p_value,
-    }
-
-
-# ============================================================
-# 5. Retrieval (Ridge LOO)
-# ============================================================
-
-def compute_retrieval_metrics(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    ks: List[int] = [1, 5, 10, 20],
-) -> Tuple[Dict[str, float], np.ndarray]:
-    """Cross-modal retrieval using Ridge LOO projection."""
-    n = module_emb.shape[0]
-    scaler_x = StandardScaler()
-    scaler_y = StandardScaler()
-    X_scaled = scaler_x.fit_transform(module_emb)
-    Y_scaled = scaler_y.fit_transform(text_emb)
-
-    ranks = np.zeros(n, dtype=np.int64)
-    for i in range(n):
-        train_mask = np.ones(n, dtype=bool)
-        train_mask[i] = False
-        ridge = RidgeCV(alphas=np.logspace(-2, 4, 7))
-        ridge.fit(X_scaled[train_mask], Y_scaled[train_mask])
-        y_pred = ridge.predict(X_scaled[i:i+1])
-        sims = cosine_similarity(y_pred, Y_scaled)[0]
-        sorted_idx = np.argsort(-sims)
-        ranks[i] = int(np.where(sorted_idx == i)[0][0])
-
-    metrics = {}
-    for k in ks:
-        metrics[f"recall@{k}"] = float(np.mean(ranks < k))
-    metrics["mrr"] = float(np.mean(1.0 / (ranks + 1)))
-    metrics["median_rank"] = int(np.median(ranks))
-    metrics["mean_rank"] = float(np.mean(ranks))
-    return metrics, ranks
-
-
-# ============================================================
-# 6. Visualization
-# ============================================================
-
-def plot_rsa_scatter(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    rsa_r: float,
-    rsa_p: float,
-    output_path: str,
-    metric: str = "cosine",
-    max_pairs: int = 5000,
-):
-    """Scatter plot of pairwise distances: module space vs text space."""
-    dist_module = pdist(module_emb, metric=metric)
-    dist_text = pdist(text_emb, metric=metric)
-
-    if len(dist_module) > max_pairs:
-        rng = np.random.default_rng(0)
-        idx = rng.choice(len(dist_module), size=max_pairs, replace=False)
-        dist_module = dist_module[idx]
-        dist_text = dist_text[idx]
-
-    fig, ax = plt.subplots(figsize=(7, 6))
-    ax.scatter(dist_text, dist_module, alpha=0.1, s=3, color="steelblue")
-    ax.set_xlabel("Text pairwise distance (cosine)")
-    ax.set_ylabel("Module embedding pairwise distance (cosine)")
-    ax.set_title(f"RSA: Spearman r = {rsa_r:.4f} (p = {rsa_p:.4f})")
-
-    z = np.polyfit(dist_text, dist_module, 1)
-    p_line = np.poly1d(z)
-    x_range = np.linspace(dist_text.min(), dist_text.max(), 100)
-    ax.plot(x_range, p_line(x_range), "r-", linewidth=2, alpha=0.8)
-
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    plt.close()
-    logger.info("Saved: %s", output_path)
-
-
-def plot_retrieval_rank_distribution(
-    ranks: np.ndarray,
-    n_total: int,
-    output_path: str,
-):
-    """Histogram of retrieval ranks."""
-    fig, ax = plt.subplots(figsize=(8, 5))
-    ax.hist(ranks, bins=50, alpha=0.7, color="darkorange", edgecolor="white")
-    ax.axvline(n_total / 2, color="gray", linestyle="--", linewidth=1.5,
-               label=f"Random baseline = {n_total//2}")
-    ax.axvline(np.median(ranks), color="blue", linestyle="-", linewidth=1.5,
-               label=f"Median = {int(np.median(ranks))}")
-    ax.set_xlabel("Retrieval Rank (0 = perfect)")
-    ax.set_ylabel("Count")
-    ax.set_title("Cross-Modal Retrieval: Rank Distribution (LOO Ridge)")
-    ax.legend()
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    plt.close()
-    logger.info("Saved: %s", output_path)
-
-
-def plot_rsa_null_distribution(
-    observed: float,
-    null_values_mean: float,
-    null_values_std: float,
-    p_value: float,
-    output_path: str,
-):
-    """Plot observed RSA vs null distribution."""
-    fig, ax = plt.subplots(figsize=(7, 5))
-    x = np.linspace(null_values_mean - 4*null_values_std,
-                    max(observed + null_values_std, null_values_mean + 4*null_values_std), 200)
-    from scipy.stats import norm as norm_dist
-    y = norm_dist.pdf(x, null_values_mean, null_values_std)
-    ax.fill_between(x, y, alpha=0.3, color="gray", label="Null distribution")
-    ax.plot(x, y, color="gray", linewidth=1.5)
-    ax.axvline(observed, color="red", linewidth=2.5, linestyle="-",
-               label=f"Observed = {observed:.4f}")
-    ax.set_xlabel("Spearman r (distance matrix correlation)")
-    ax.set_ylabel("Density")
-    ax.set_title(f"Mantel Test: p = {p_value:.4f}")
-    ax.legend()
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    plt.close()
-    logger.info("Saved: %s", output_path)
-
-
-def save_retrieval_examples(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    texts: List[str],
-    names: List[str],
-    output_path: str,
-    top_k: int = 5,
-    query_indices: Optional[List[int]] = None,
-):
-    """Save retrieval examples for selected modules."""
-    n = module_emb.shape[0]
-    scaler_x = StandardScaler()
-    scaler_y = StandardScaler()
-    X = scaler_x.fit_transform(module_emb)
-    Y = scaler_y.fit_transform(text_emb)
-
-    ridge = RidgeCV(alphas=np.logspace(-2, 4, 7))
-    ridge.fit(X, Y)
-
-    if query_indices is None:
-        rng = np.random.default_rng(0)
-        query_indices = sorted(rng.choice(n, size=min(20, n), replace=False))
-
-    rows = []
-    for qi in query_indices:
-        y_pred = ridge.predict(X[qi:qi+1])
-        sims = cosine_similarity(y_pred, Y)[0]
-        top_idx = np.argsort(-sims)[:top_k]
-        for rank, ti in enumerate(top_idx):
-            rows.append({
-                "query_module": names[qi],
-                "query_text": texts[qi],
-                "retrieved_rank": rank + 1,
-                "retrieved_text": texts[ti],
-                "similarity": float(sims[ti]),
-                "is_self": ti == qi,
-            })
-
-    df = pd.DataFrame(rows)
-    df.to_csv(output_path, index=False)
-    logger.info("Saved retrieval examples: %s", output_path)
-
-
-# ============================================================
-# 7. Main
-# ============================================================
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description=__doc__,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    parser.add_argument("--biomod_embedding_path", "--module_embedding_path", type=str, required=True,
-                        help="Path to module_embeddings.npy from experiment 1.", dest='module_embedding_path')
-    parser.add_argument("--biomod_json", "--module_json", type=str,
-                        default=str(REPO_ROOT / "model_dependencies" / "biomod_data_all" / "merged_biomods.json"),
-                        help="Path to merged_biomods.json.", dest='module_json')
-    parser.add_argument("--descriptions_file", type=str, default=None,
-                        help="Optional JSON with extended pathway descriptions.")
-    parser.add_argument("--sentence_model", type=str,
-                        default="pritamdeka/S-PubMedBert-MS-MARCO",
-                        help="HuggingFace sentence-transformers model.")
-    parser.add_argument("--device", type=str, default="cpu",
-                        help="Device for sentence encoder.")
-    parser.add_argument("--n_permutations", type=int, default=1000)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--output_dir", type=str, required=True)
-    parser.add_argument("--skip_loo", action="store_true",
-                        help="Skip LOO retrieval (slow). Use RSA/CKA only.")
-    parser.add_argument("--rsa_metric", type=str, default="cosine",
-                        choices=["cosine", "euclidean"],
-                        help="Distance metric for RSA.")
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s [%(levelname)s] %(message)s",
-        handlers=[
-            logging.StreamHandler(),
-            logging.FileHandler(output_dir / "run.log", mode="w"),
-        ],
-    )
-
-    # --- Load module names and embeddings ---
-    logger.info("Loading module names from %s", args.module_json)
-    module_names = load_biomod_names(args.module_json)
-    logger.info("Total modules: %d", len(module_names))
-
-    logger.info("Loading module embeddings from %s", args.module_embedding_path)
-    all_embeddings = np.load(args.module_embedding_path)
-    assert all_embeddings.shape[0] == len(module_names), (
-        f"Embedding count {all_embeddings.shape[0]} != module count {len(module_names)}"
-    )
-
-    # --- Build text descriptions (Hallmarks + Reactome only) ---
-    indices, names, texts = build_text_descriptions(
-        module_names, args.descriptions_file
-    )
-    logger.info("Modules with text descriptions: %d / %d", len(indices), len(module_names))
-
-    module_emb = all_embeddings[indices]
-    logger.info("Module embedding shape: %s", module_emb.shape)
-
-    # --- Encode texts ---
-    text_emb = encode_texts(
-        texts, model_name=args.sentence_model, device=args.device
-    )
-    logger.info("Text embedding shape: %s", text_emb.shape)
-
-    # Save intermediate
-    np.save(output_dir / "text_embeddings.npy", text_emb)
-    np.save(output_dir / "module_embeddings_subset.npy", module_emb)
-    with open(output_dir / "pathway_descriptions.json", "w", encoding="utf-8") as f:
-        json.dump(
-            [{"index": idx, "module_name": name, "text": text}
-             for idx, name, text in zip(indices, names, texts)],
-            f, indent=2, ensure_ascii=False,
-        )
-
-    # --- RSA / Mantel test (PRIMARY METRIC) ---
-    logger.info("=" * 50)
-    logger.info("RSA / Mantel Test (primary alignment metric)")
-    logger.info("=" * 50)
-
-    sp_r, sp_p, pe_r, pe_p = compute_rsa(module_emb, text_emb, metric=args.rsa_metric)
-    logger.info("RSA (all modules): Spearman r=%.4f (p=%.2e), Pearson r=%.4f (p=%.2e)",
-                sp_r, sp_p, pe_r, pe_p)
-
-    mantel_results = mantel_test(
-        module_emb, text_emb, metric=args.rsa_metric,
-        n_perm=args.n_permutations, seed=args.seed,
-    )
-    logger.info("Mantel test: observed=%.4f, null=%.4f +/- %.4f, p=%.4f, z=%.1f",
-                mantel_results["rsa_spearman_observed"],
-                mantel_results["rsa_null_mean"],
-                mantel_results["rsa_null_std"],
-                mantel_results["rsa_p_value"],
-                mantel_results["rsa_effect_size_z"])
-
-    # RSA by family
-    family_rsa = compute_rsa_by_family(
-        module_emb, text_emb, names, metric=args.rsa_metric,
-        n_perm=min(args.n_permutations, 500), seed=args.seed,
-    )
-    for fam, res in family_rsa.items():
-        logger.info("RSA [%s]: r=%.4f, null=%.4f +/- %.4f, p=%.4f, z=%.1f",
-                    fam, res["rsa_spearman_observed"], res["rsa_null_mean"],
-                    res["rsa_null_std"], res["rsa_p_value"], res["rsa_effect_size_z"])
-
-    # --- CKA ---
-    logger.info("Computing CKA with %d permutations...", args.n_permutations)
-    cka_results = permutation_test_cka(
-        module_emb, text_emb, n_perm=args.n_permutations, seed=args.seed
-    )
-    logger.info("CKA: observed=%.4f, null=%.4f +/- %.4f, p=%.4f",
-                cka_results["cka_observed"], cka_results["cka_null_mean"],
-                cka_results["cka_null_std"], cka_results["cka_p_value"])
-
-    # --- LOO retrieval (optional) ---
-    retrieval_results = {}
-    ranks = None
-    if not args.skip_loo:
-        logger.info("Computing LOO retrieval metrics (n=%d)...", len(indices))
-        retrieval_results, ranks = compute_retrieval_metrics(module_emb, text_emb)
-        logger.info("Retrieval: recall@1=%.3f, @5=%.3f, @10=%.3f, @20=%.3f, MRR=%.3f, median_rank=%d",
-                    retrieval_results.get("recall@1", 0),
-                    retrieval_results.get("recall@5", 0),
-                    retrieval_results.get("recall@10", 0),
-                    retrieval_results.get("recall@20", 0),
-                    retrieval_results.get("mrr", 0),
-                    retrieval_results.get("median_rank", 0))
-
-    # --- Aggregate metrics ---
-    metrics = {
-        "n_modules_with_text": len(indices),
-        "n_modules_total": len(module_names),
-        "module_emb_dim": int(module_emb.shape[1]),
-        "text_emb_dim": int(text_emb.shape[1]),
-        "sentence_model": args.sentence_model,
-        "rsa_metric": args.rsa_metric,
-        "n_permutations": args.n_permutations,
-        "rsa_spearman_r": sp_r,
-        "rsa_spearman_p": sp_p,
-        "rsa_pearson_r": pe_r,
-        "rsa_pearson_p": pe_p,
-        **mantel_results,
-        "rsa_by_family": family_rsa,
-        **cka_results,
-        **retrieval_results,
-    }
-
-    with open(output_dir / "alignment_metrics.json", "w") as f:
-        json.dump(metrics, f, indent=2)
-    logger.info("Saved alignment_metrics.json")
-
-    # --- Plots ---
-    plot_rsa_scatter(
-        module_emb, text_emb, rsa_r=sp_r, rsa_p=mantel_results["rsa_p_value"],
-        output_path=str(output_dir / "rsa_scatter.png"),
-        metric=args.rsa_metric,
-    )
-
-    plot_rsa_null_distribution(
-        observed=mantel_results["rsa_spearman_observed"],
-        null_values_mean=mantel_results["rsa_null_mean"],
-        null_values_std=mantel_results["rsa_null_std"],
-        p_value=mantel_results["rsa_p_value"],
-        output_path=str(output_dir / "rsa_null_distribution.png"),
-    )
-
-    if ranks is not None:
-        plot_retrieval_rank_distribution(
-            ranks, n_total=len(indices),
-            output_path=str(output_dir / "retrieval_ranks.png"),
-        )
-
-    # --- Retrieval examples ---
-    representative_keywords = [
-        "interferon", "cell cycle", "oxidative", "apoptosis",
-        "immune", "signaling", "metabolism", "dna repair",
-        "nfkb", "mtor", "wnt", "hedgehog", "notch",
-        "interleukin", "toll", "mapk",
-    ]
-    query_indices = []
-    for kw in representative_keywords:
-        for i, t in enumerate(texts):
-            if kw.lower() in t.lower() and i not in query_indices:
-                query_indices.append(i)
-                break
-    if len(query_indices) < 10:
-        rng = np.random.default_rng(args.seed)
-        extra = rng.choice(len(texts), size=15, replace=False)
-        query_indices = list(set(query_indices) | set(extra.tolist()))[:20]
-
-    save_retrieval_examples(
-        module_emb, text_emb, texts, names,
-        output_path=str(output_dir / "retrieval_examples.csv"),
-        query_indices=sorted(query_indices),
-    )
-
-    # --- Summary ---
-    logger.info("=" * 60)
-    logger.info("EXPERIMENT 2 SUMMARY")
-    logger.info("=" * 60)
-    logger.info("Modules with text: %d / %d", len(indices), len(module_names))
-    logger.info("RSA (Mantel): Spearman r=%.4f, p=%.4f, z=%.1f sigma",
-                mantel_results["rsa_spearman_observed"],
-                mantel_results["rsa_p_value"],
-                mantel_results["rsa_effect_size_z"])
-    logger.info("CKA: %.4f (p=%.4f)", cka_results["cka_observed"], cka_results["cka_p_value"])
-    if retrieval_results:
-        logger.info("Retrieval: recall@1=%.3f, @5=%.3f, @10=%.3f, MRR=%.3f, median_rank=%d",
-                    retrieval_results["recall@1"], retrieval_results["recall@5"],
-                    retrieval_results["recall@10"], retrieval_results["mrr"],
-                    retrieval_results["median_rank"])
-    logger.info("Output: %s", output_dir)
-    logger.info("=" * 60)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_extended.py

@@ -1,537 +0,0 @@
-"""Extended analysis for Experiment 2 — Module-Text Alignment.
-
-Three analyses:
-1. Reactome category-level RSA (aggregate 491 pathways into ~16 top-level categories)
-2. Reactome with official descriptions (via descriptions JSON)
-3. Hallmarks deep dive (per-module alignment scores, heatmap)
-
-Usage (run from the BioMod repository root):
-    python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_extended.py \
-        --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-        --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-        --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-        --output_dir outputs/interpretability/text_alignment/pretrain_subset_v2_extended
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-from collections import Counter, defaultdict
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-from scipy.spatial.distance import pdist, squareform
-from scipy.stats import spearmanr
-from sklearn.linear_model import RidgeCV
-from sklearn.metrics.pairwise import cosine_similarity
-from sklearn.preprocessing import StandardScaler
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-HERE = Path(__file__).resolve().parent
-INTERP_DIR = HERE.parent
-BIOMOD_SOURCE = INTERP_DIR.parent
-REPO_ROOT = BIOMOD_SOURCE.parent
-for p in (str(INTERP_DIR), str(BIOMOD_SOURCE), str(REPO_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from interpretability.utils import load_biomod_names, parse_biomod_family
-from biomod_text_alignment.biomod_text_alignment import (
-    build_text_descriptions,
-    compute_cka,
-    compute_rsa,
-    encode_texts,
-    mantel_test,
-    parse_pathway_text,
-)
-
-logger = logging.getLogger(__name__)
-
-
-# ============================================================
-# Reactome top-level category mapping
-# ============================================================
-
-REACTOME_CATEGORIES = {
-    "Cell Cycle": ["cell cycle", "mitotic", "mitosis", "meiosis", "g1/s", "g2/m",
-                   "s phase", "m phase", "checkpoint", "cdc", "apc/c", "cyclin",
-                   "condensation of", "centrosome", "kinetochore", "spindle",
-                   "cohes", "pre-replicative", "separation of sister", "aurora"],
-    "DNA Repair": ["dna repair", "base excision", "nucleotide excision",
-                   "mismatch repair", "homologous recombination", "non-homologous",
-                   "double-strand", "single-strand", "fanconi", "dna damage"],
-    "DNA Replication": ["dna replication", "replication", "okazaki", "telomere"],
-    "Signal Transduction": ["signaling", "signal transduction", "mapk", "ras",
-                            "raf", "mek", "erk", "pi3k", "akt", "mtor", "wnt",
-                            "notch", "hedgehog", "tgf-beta", "bmp", "fgfr",
-                            "egfr", "vegf", "pdgf", "igf", "insulin", "erbb",
-                            "met receptor", "receptor tyrosine", "gpcr",
-                            "g alpha", "rho gtpase", "rac1", "cdc42",
-                            "ca2+ pathway", "nfkb", "nf-kb", "nmda",
-                            "calcium", "pkc", "plc", "dag", "pip"],
-    "Immune System": ["immune", "innate", "adaptive", "interferon", "interleukin",
-                      "cytokine", "complement", "toll", "antigen", "mhc", "tcr",
-                      "bcr", "fcr", "neutrophil", "phagocyt", "inflamm",
-                      "antimicrobial", "antiviral", "ifn-stimulated",
-                      "c-type lectin", "nod-like", "rig-i"],
-    "Metabolism": ["metabolism", "metabolic", "glycolysis", "gluconeogenesis",
-                   "tca cycle", "citric acid", "oxidative phosphorylation",
-                   "fatty acid", "lipid", "cholesterol", "bile", "amino acid",
-                   "nucleotide", "purine", "pyrimidine", "pentose", "urea",
-                   "biological oxidation", "glycosylation", "biosynthesis",
-                   "synthesis of", "respiratory electron"],
-    "Gene Expression": ["transcription", "translation", "mrna", "rrna", "trna",
-                        "splicing", "polyadenylation", "ribosom",
-                        "eukaryotic translation", "nonsense-mediated",
-                        "deadenylation", "rna polymerase", "processing of capped"],
-    "Protein Homeostasis": ["ubiquitin", "proteasome", "sumo", "nedd",
-                            "autophagy", "chaperone", "unfolded protein",
-                            "protein folding", "hsp", "degradation of",
-                            "deubiquitin"],
-    "Transport": ["transport", "vesicle", "endocyt", "exocyt", "golgi",
-                  "er to golgi", "membrane trafficking", "clathrin", "copi",
-                  "copii", "ion channel", "slc-mediated", "abc transporter",
-                  "nuclear pore", "nuclear import", "nuclear export", "cargo"],
-    "Apoptosis": ["apoptosis", "apoptotic", "programmed cell death", "caspase",
-                  "death receptor", "bcl-2"],
-    "Cell-Cell Communication": ["cell junction", "gap junction", "tight junction",
-                                "adherens", "cell-cell", "extracellular matrix",
-                                "integrin", "synapse", "neurotransmitter",
-                                "collagen", "laminin", "elastic fibre",
-                                "anchoring"],
-    "Chromatin & Epigenetics": ["chromatin", "histone", "methylation",
-                                "acetylation", "deacetylation", "epigenetic",
-                                "polycomb", "hox"],
-    "Hemostasis": ["hemostasis", "platelet", "coagulation", "fibrin", "thrombin"],
-    "Developmental Biology": ["developmental", "axon", "neural", "muscle",
-                              "myogenesis", "keratinization", "differentiation",
-                              "cardiac conduction"],
-    "Disease": ["disease", "hiv", "influenza", "sars", "hcmv", "herpes",
-                "hepatitis", "cancer", "oncogene", "tumor", "amyloid"],
-    "Cellular Responses": ["cellular response", "oxidative stress", "hypoxia",
-                           "heat shock", "senescence", "ferroptosis",
-                           "detoxification"],
-}
-
-
-def categorize_reactome(pathway_name: str) -> str:
-    name_lower = pathway_name.lower()
-    for cat, keywords in REACTOME_CATEGORIES.items():
-        for kw in keywords:
-            if kw in name_lower:
-                return cat
-    return "Other"
-
-
-# ============================================================
-# Analysis 1: Reactome category-level RSA
-# ============================================================
-
-def reactome_category_rsa(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    output_dir: Path,
-    sentence_model: str,
-    device: str,
-    n_perm: int = 1000,
-    seed: int = 42,
-) -> Dict:
-    """Aggregate Reactome modules by top-level category, then do RSA."""
-    logger.info("=" * 50)
-    logger.info("Analysis 1: Reactome category-level RSA")
-    logger.info("=" * 50)
-
-    families = [parse_biomod_family(n) for n in names]
-    reactome_mask = np.array([f == "Reactome" for f in families])
-    reactome_idx = np.where(reactome_mask)[0]
-
-    reactome_names = [names[i] for i in reactome_idx]
-    reactome_texts = [texts[i] for i in reactome_idx]
-    reactome_module_emb = module_emb[reactome_idx]
-
-    cats = [categorize_reactome(t) for t in reactome_texts]
-    cat_counts = Counter(cats)
-    logger.info("Reactome category distribution:")
-    for cat, count in sorted(cat_counts.items(), key=lambda x: -x[1]):
-        logger.info("  %s: %d", cat, count)
-
-    # Aggregate: mean embedding per category
-    unique_cats = sorted(set(c for c in cats if c != "Other"))
-    n_cats = len(unique_cats)
-    logger.info("Categories (excl. Other): %d", n_cats)
-
-    cat_module_emb = np.zeros((n_cats, module_emb.shape[1]))
-    for i, cat in enumerate(unique_cats):
-        mask = np.array([c == cat for c in cats])
-        cat_module_emb[i] = reactome_module_emb[mask].mean(axis=0)
-
-    # Encode category names as text
-    cat_text_emb = encode_texts(unique_cats, model_name=sentence_model, device=device)
-
-    # RSA on category-level
-    sp_r, sp_p, pe_r, pe_p = compute_rsa(cat_module_emb, cat_text_emb, metric="cosine")
-    logger.info("Category-level RSA: Spearman r=%.4f (p=%.4f), Pearson r=%.4f",
-                sp_r, sp_p, pe_r)
-
-    mantel = mantel_test(cat_module_emb, cat_text_emb, metric="cosine",
-                         n_perm=n_perm, seed=seed)
-    logger.info("Mantel test: observed=%.4f, null=%.4f +/- %.4f, p=%.4f, z=%.1f",
-                mantel["rsa_spearman_observed"], mantel["rsa_null_mean"],
-                mantel["rsa_null_std"], mantel["rsa_p_value"],
-                mantel["rsa_effect_size_z"])
-
-    cka = compute_cka(cat_module_emb, cat_text_emb)
-    logger.info("Category-level CKA: %.4f", cka)
-
-    # Visualize: distance heatmaps
-    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
-
-    dist_module = squareform(pdist(cat_module_emb, metric="cosine"))
-    dist_text = squareform(pdist(cat_text_emb, metric="cosine"))
-
-    im0 = axes[0].imshow(dist_module, cmap="viridis", aspect="auto")
-    axes[0].set_xticks(range(n_cats))
-    axes[0].set_xticklabels(unique_cats, rotation=45, ha="right", fontsize=7)
-    axes[0].set_yticks(range(n_cats))
-    axes[0].set_yticklabels(unique_cats, fontsize=7)
-    axes[0].set_title("Module Embedding Distance (cosine)")
-    plt.colorbar(im0, ax=axes[0], fraction=0.046)
-
-    im1 = axes[1].imshow(dist_text, cmap="viridis", aspect="auto")
-    axes[1].set_xticks(range(n_cats))
-    axes[1].set_xticklabels(unique_cats, rotation=45, ha="right", fontsize=7)
-    axes[1].set_yticks(range(n_cats))
-    axes[1].set_yticklabels(unique_cats, fontsize=7)
-    axes[1].set_title("Text Embedding Distance (cosine)")
-    plt.colorbar(im1, ax=axes[1], fraction=0.046)
-
-    plt.suptitle(f"Reactome Category-Level: RSA r={sp_r:.3f}, p={mantel['rsa_p_value']:.4f}")
-    plt.tight_layout()
-    plt.savefig(output_dir / "reactome_category_heatmaps.png", dpi=150, bbox_inches="tight")
-    plt.close()
-
-    results = {
-        "n_categories": n_cats,
-        "category_counts": dict(cat_counts),
-        "rsa_spearman_r": sp_r,
-        "rsa_spearman_p": sp_p,
-        "rsa_pearson_r": pe_r,
-        "mantel": mantel,
-        "cka": cka,
-        "categories": unique_cats,
-    }
-    return results
-
-
-# ============================================================
-# Analysis 2: Reactome with descriptions
-# ============================================================
-
-def reactome_with_descriptions(
-    module_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    descriptions_file: str,
-    output_dir: Path,
-    sentence_model: str,
-    device: str,
-    n_perm: int = 1000,
-    seed: int = 42,
-) -> Optional[Dict]:
-    """Re-run RSA using extended descriptions for Reactome pathways."""
-    if not descriptions_file or not Path(descriptions_file).exists():
-        logger.info("No descriptions file provided, skipping Analysis 2.")
-        return None
-
-    logger.info("=" * 50)
-    logger.info("Analysis 2: Reactome with official descriptions")
-    logger.info("=" * 50)
-
-    with open(descriptions_file, "r", encoding="utf-8") as f:
-        descriptions = json.load(f)
-
-    families = [parse_biomod_family(n) for n in names]
-    reactome_mask = np.array([f == "Reactome" for f in families])
-    reactome_idx = np.where(reactome_mask)[0]
-
-    # Build enriched texts
-    enriched_texts = []
-    valid_idx = []
-    for i in reactome_idx:
-        pname = names[i]
-        base_text = texts[i]
-        if pname in descriptions and descriptions[pname]:
-            enriched_texts.append(f"{base_text}: {descriptions[pname]}")
-            valid_idx.append(i)
-        else:
-            enriched_texts.append(base_text)
-            valid_idx.append(i)
-
-    logger.info("Reactome modules with descriptions: %d / %d",
-                sum(1 for i in reactome_idx if names[i] in descriptions),
-                len(reactome_idx))
-
-    sub_module_emb = module_emb[valid_idx]
-    text_emb_enriched = encode_texts(enriched_texts, model_name=sentence_model, device=device)
-
-    sp_r, sp_p, pe_r, pe_p = compute_rsa(sub_module_emb, text_emb_enriched, metric="cosine")
-    logger.info("RSA with descriptions: Spearman r=%.4f (p=%.2e)", sp_r, sp_p)
-
-    mantel = mantel_test(sub_module_emb, text_emb_enriched, metric="cosine",
-                         n_perm=n_perm, seed=seed)
-    logger.info("Mantel: observed=%.4f, null=%.4f +/- %.4f, p=%.4f, z=%.1f",
-                mantel["rsa_spearman_observed"], mantel["rsa_null_mean"],
-                mantel["rsa_null_std"], mantel["rsa_p_value"],
-                mantel["rsa_effect_size_z"])
-
-    return {
-        "rsa_spearman_r": sp_r,
-        "rsa_spearman_p": sp_p,
-        "mantel": mantel,
-    }
-
-
-# ============================================================
-# Analysis 3: Hallmarks deep dive
-# ============================================================
-
-def hallmarks_deep_dive(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    output_dir: Path,
-    seed: int = 42,
-) -> Dict:
-    """Per-module alignment analysis for the 50 Hallmarks."""
-    logger.info("=" * 50)
-    logger.info("Analysis 3: Hallmarks deep dive")
-    logger.info("=" * 50)
-
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    h_texts = [texts[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    logger.info("Hallmarks: %d modules", n_h)
-
-    # Per-module alignment: for each hallmark, compute cosine sim between
-    # its module embedding's nearest neighbors and its text embedding's nearest neighbors
-    # (neighborhood overlap)
-    module_sim_matrix = cosine_similarity(h_module)
-    text_sim_matrix = cosine_similarity(h_text)
-
-    # Per-module: correlation between its similarity profile in module space vs text space
-    per_module_corr = np.zeros(n_h)
-    for i in range(n_h):
-        mask = np.ones(n_h, dtype=bool)
-        mask[i] = False
-        r, _ = spearmanr(module_sim_matrix[i, mask], text_sim_matrix[i, mask])
-        per_module_corr[i] = r
-
-    # Sort by alignment quality
-    sort_idx = np.argsort(-per_module_corr)
-    logger.info("\nTop-10 best aligned Hallmarks:")
-    for rank, i in enumerate(sort_idx[:10]):
-        logger.info("  %2d. %s (r=%.3f)", rank+1, h_texts[i], per_module_corr[i])
-
-    logger.info("\nBottom-10 worst aligned Hallmarks:")
-    for rank, i in enumerate(sort_idx[-10:]):
-        logger.info("  %2d. %s (r=%.3f)", rank+1, h_texts[i], per_module_corr[i])
-
-    # Save per-module scores
-    df = pd.DataFrame({
-        "module_name": h_names,
-        "text": h_texts,
-        "alignment_r": per_module_corr,
-    }).sort_values("alignment_r", ascending=False)
-    df.to_csv(output_dir / "hallmarks_per_module_alignment.csv", index=False)
-
-    # Heatmap: module similarity vs text similarity
-    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-
-    im0 = axes[0].imshow(module_sim_matrix[sort_idx][:, sort_idx], cmap="RdBu_r",
-                          vmin=-0.5, vmax=0.5, aspect="auto")
-    axes[0].set_title("Module Embedding Similarity")
-    axes[0].set_xticks(range(n_h))
-    axes[0].set_xticklabels([h_texts[i][:15] for i in sort_idx], rotation=90, fontsize=5)
-    axes[0].set_yticks(range(n_h))
-    axes[0].set_yticklabels([h_texts[i][:15] for i in sort_idx], fontsize=5)
-    plt.colorbar(im0, ax=axes[0], fraction=0.046)
-
-    im1 = axes[1].imshow(text_sim_matrix[sort_idx][:, sort_idx], cmap="RdBu_r",
-                          vmin=-0.5, vmax=0.5, aspect="auto")
-    axes[1].set_title("Text Embedding Similarity")
-    axes[1].set_xticks(range(n_h))
-    axes[1].set_xticklabels([h_texts[i][:15] for i in sort_idx], rotation=90, fontsize=5)
-    axes[1].set_yticks([])
-    plt.colorbar(im1, ax=axes[1], fraction=0.046)
-
-    # Bar chart of per-module alignment
-    colors = ["steelblue" if r > 0 else "salmon" for r in per_module_corr[sort_idx]]
-    axes[2].barh(range(n_h), per_module_corr[sort_idx], color=colors)
-    axes[2].set_yticks(range(n_h))
-    axes[2].set_yticklabels([h_texts[i][:20] for i in sort_idx], fontsize=5)
-    axes[2].set_xlabel("Spearman r (neighborhood alignment)")
-    axes[2].set_title("Per-Hallmark Alignment Score")
-    axes[2].axvline(0, color="black", linewidth=0.5)
-
-    plt.tight_layout()
-    plt.savefig(output_dir / "hallmarks_deep_dive.png", dpi=150, bbox_inches="tight")
-    plt.close()
-    logger.info("Saved hallmarks_deep_dive.png")
-
-    # kNN overlap analysis
-    k = 5
-    module_knn = np.argsort(-module_sim_matrix, axis=1)[:, 1:k+1]
-    text_knn = np.argsort(-text_sim_matrix, axis=1)[:, 1:k+1]
-    knn_overlap = np.zeros(n_h)
-    for i in range(n_h):
-        overlap = len(set(module_knn[i]) & set(text_knn[i]))
-        knn_overlap[i] = overlap / k
-
-    logger.info("\nkNN overlap (k=%d): mean=%.3f, median=%.3f",
-                k, knn_overlap.mean(), np.median(knn_overlap))
-
-    # Permutation test for kNN overlap
-    rng = np.random.default_rng(seed)
-    n_perm = 1000
-    null_overlaps = np.zeros(n_perm)
-    for p in range(n_perm):
-        perm = rng.permutation(n_h)
-        text_knn_perm = text_knn[perm]
-        overlaps_perm = np.zeros(n_h)
-        for i in range(n_h):
-            overlaps_perm[i] = len(set(module_knn[i]) & set(text_knn_perm[i])) / k
-        null_overlaps[p] = overlaps_perm.mean()
-
-    knn_p = float(np.mean(null_overlaps >= knn_overlap.mean()))
-    logger.info("kNN overlap permutation test: observed=%.3f, null=%.3f +/- %.3f, p=%.4f",
-                knn_overlap.mean(), null_overlaps.mean(), null_overlaps.std(), knn_p)
-
-    results = {
-        "n_hallmarks": n_h,
-        "per_module_alignment_mean": float(per_module_corr.mean()),
-        "per_module_alignment_median": float(np.median(per_module_corr)),
-        "per_module_alignment_std": float(per_module_corr.std()),
-        "n_positive_alignment": int((per_module_corr > 0).sum()),
-        "n_negative_alignment": int((per_module_corr < 0).sum()),
-        "knn_k": k,
-        "knn_overlap_mean": float(knn_overlap.mean()),
-        "knn_overlap_null_mean": float(null_overlaps.mean()),
-        "knn_overlap_p": knn_p,
-    }
-    return results
-
-
-# ============================================================
-# Main
-# ============================================================
-
-def parse_args():
-    parser = argparse.ArgumentParser(description=__doc__,
-                                     formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument("--biomod_embedding_path", "--module_embedding_path", type=str, required=True, dest='module_embedding_path')
-    parser.add_argument("--biomod_json", "--module_json", type=str,
-                        default=str(REPO_ROOT / "model_dependencies" / "biomod_data_all" / "merged_biomods.json"), dest='module_json')
-    parser.add_argument("--descriptions_file", type=str, default=None,
-                        help="JSON with Reactome summaryText descriptions.")
-    parser.add_argument("--sentence_model", type=str,
-                        default="pritamdeka/S-PubMedBert-MS-MARCO")
-    parser.add_argument("--device", type=str, default="cpu")
-    parser.add_argument("--n_permutations", type=int, default=1000)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--output_dir", type=str, required=True)
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s [%(levelname)s] %(message)s",
-        handlers=[
-            logging.StreamHandler(),
-            logging.FileHandler(output_dir / "run_extended.log", mode="w"),
-        ],
-    )
-
-    # Load data
-    module_names = load_biomod_names(args.module_json)
-    all_embeddings = np.load(args.module_embedding_path)
-    indices, names, texts = build_text_descriptions(module_names)
-    module_emb = all_embeddings[indices]
-
-    # Load or compute text embeddings
-    text_emb_path = output_dir.parent / "pretrain_subset_v2" / "text_embeddings.npy"
-    if text_emb_path.exists():
-        logger.info("Loading cached text embeddings from %s", text_emb_path)
-        text_emb = np.load(text_emb_path)
-    else:
-        text_emb = encode_texts(texts, model_name=args.sentence_model, device=args.device)
-
-    all_results = {}
-
-    # Analysis 1: Reactome category-level
-    cat_results = reactome_category_rsa(
-        module_emb, text_emb, names, texts, output_dir,
-        sentence_model=args.sentence_model, device=args.device,
-        n_perm=args.n_permutations, seed=args.seed,
-    )
-    all_results["reactome_category_rsa"] = cat_results
-
-    # Analysis 2: Reactome with descriptions
-    desc_results = reactome_with_descriptions(
-        module_emb, names, texts, args.descriptions_file, output_dir,
-        sentence_model=args.sentence_model, device=args.device,
-        n_perm=args.n_permutations, seed=args.seed,
-    )
-    if desc_results:
-        all_results["reactome_with_descriptions"] = desc_results
-
-    # Analysis 3: Hallmarks deep dive
-    hallmark_results = hallmarks_deep_dive(
-        module_emb, text_emb, names, texts, output_dir, seed=args.seed,
-    )
-    all_results["hallmarks_deep_dive"] = hallmark_results
-
-    # Save all results
-    with open(output_dir / "extended_metrics.json", "w") as f:
-        json.dump(all_results, f, indent=2, default=str)
-    logger.info("Saved extended_metrics.json")
-
-    # Summary
-    logger.info("=" * 60)
-    logger.info("EXTENDED ANALYSIS SUMMARY")
-    logger.info("=" * 60)
-    logger.info("1. Reactome category-level RSA: r=%.4f, p=%.4f",
-                cat_results["rsa_spearman_r"], cat_results["mantel"]["rsa_p_value"])
-    logger.info("2. Hallmarks per-module alignment: mean r=%.3f, %d/%d positive",
-                hallmark_results["per_module_alignment_mean"],
-                hallmark_results["n_positive_alignment"],
-                hallmark_results["n_hallmarks"])
-    logger.info("3. Hallmarks kNN overlap: %.3f (null=%.3f, p=%.4f)",
-                hallmark_results["knn_overlap_mean"],
-                hallmark_results["knn_overlap_null_mean"],
-                hallmark_results["knn_overlap_p"])
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_publication.py

@@ -1,1277 +0,0 @@
-"""Publication-quality analysis for Experiment 2 — Module-Text Alignment.
-
-Produces 6 main figures + comprehensive metrics with bootstrap CIs for paper.
-Supports multi-model comparison and cross-dataset validation.
-
-Usage (run from the BioMod repository root):
-    python BioMod/interpretability/biomod_text_alignment/biomod_text_alignment_publication.py \
-        --biomod_embedding_path outputs/interpretability/geometry/pretrain_subset/module_embeddings.npy \
-        --biomod_embedding_path2 outputs/interpretability/geometry/pancreas_scib/module_embeddings.npy \
-        --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-        --sentence_model pritamdeka/S-PubMedBert-MS-MARCO \
-        --sentence_models_compare "sentence-transformers/all-mpnet-base-v2" \
-        --descriptions_file BioMod/interpretability/biomod_text_alignment/reactome_descriptions.json \
-        --n_permutations 1000 \
-        --n_bootstrap 2000 \
-        --output_dir outputs/interpretability/text_alignment/publication
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-from collections import Counter, defaultdict
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-from scipy.cluster.hierarchy import dendrogram, fcluster, linkage
-from scipy.spatial.distance import pdist, squareform
-from scipy.stats import kruskal, spearmanr
-from sklearn.metrics.pairwise import cosine_similarity
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-from matplotlib.patches import Patch as LegendHandle
-from matplotlib.lines import Line2D
-
-HERE = Path(__file__).resolve().parent
-INTERP_DIR = HERE.parent
-BIOMOD_SOURCE = INTERP_DIR.parent
-REPO_ROOT = BIOMOD_SOURCE.parent
-for p in (str(HERE), str(INTERP_DIR), str(BIOMOD_SOURCE), str(REPO_ROOT)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from interpretability.utils import (
-    HALLMARK_PROCESS,
-    load_biomod_names,
-    parse_hallmark_process,
-    parse_biomod_family,
-)
-from biomod_text_alignment.biomod_text_alignment import (
-    build_text_descriptions,
-    compute_cka,
-    compute_rsa,
-    encode_texts,
-    mantel_test,
-    permutation_test_cka,
-)
-from biomod_text_alignment.biomod_text_alignment_extended import (
-    REACTOME_CATEGORIES,
-    categorize_reactome,
-)
-
-logger = logging.getLogger(__name__)
-
-# ============================================================
-# Color palette for MSigDB hallmark process categories
-# ============================================================
-
-PROCESS_COLORS = {
-    "immune": "#2196F3",
-    "metabolic": "#4CAF50",
-    "signaling": "#FF9800",
-    "proliferation": "#E91E63",
-    "development": "#9C27B0",
-    "dna_damage": "#F44336",
-    "pathway": "#607D8B",
-    "cellular_component": "#795548",
-}
-
-PROCESS_LABELS = {
-    "immune": "Immune",
-    "metabolic": "Metabolic",
-    "signaling": "Signaling",
-    "proliferation": "Proliferation",
-    "development": "Development",
-    "dna_damage": "DNA Damage",
-    "pathway": "Pathway",
-    "cellular_component": "Cellular Component",
-}
-
-# Reactome super-categories (merge related categories for granularity sweep)
-REACTOME_SUPER_CATEGORIES = {
-    "Information Processing": ["Gene Expression", "DNA Replication", "DNA Repair",
-                               "Chromatin & Epigenetics"],
-    "Signal & Communication": ["Signal Transduction", "Cell-Cell Communication"],
-    "Cell Fate": ["Cell Cycle", "Apoptosis", "Cellular Responses"],
-    "Metabolism & Transport": ["Metabolism", "Transport"],
-    "Immune & Defense": ["Immune System", "Hemostasis", "Disease"],
-    "Development": ["Developmental Biology", "Protein Homeostasis"],
-}
-
-
-# ============================================================
-# Utility functions
-# ============================================================
-
-def savefig(fig, path_stem: str, dpi: int = 300):
-    """Save figure as both PDF and PNG."""
-    fig.savefig(f"{path_stem}.pdf", dpi=dpi, bbox_inches="tight")
-    fig.savefig(f"{path_stem}.png", dpi=dpi, bbox_inches="tight")
-    plt.close(fig)
-    logger.info("Saved: %s.{pdf,png}", path_stem)
-
-
-def adaptive_vlim(matrix: np.ndarray, percentile: int = 98, symmetric: bool = True):
-    """Compute adaptive vmin/vmax from data percentiles."""
-    vals = matrix[np.triu_indices_from(matrix, k=1)]
-    if symmetric:
-        vlim = np.percentile(np.abs(vals), percentile)
-        return -vlim, vlim
-    else:
-        return np.percentile(vals, [100 - percentile, percentile])
-
-
-def bootstrap_ci(
-    values: np.ndarray,
-    stat_fn=np.mean,
-    n_boot: int = 2000,
-    ci: float = 0.95,
-    seed: int = 42,
-) -> Tuple[float, float, float]:
-    """Bootstrap confidence interval. Returns (estimate, ci_low, ci_high)."""
-    rng = np.random.default_rng(seed)
-    n = len(values)
-    boot_stats = np.zeros(n_boot)
-    for i in range(n_boot):
-        idx = rng.integers(0, n, size=n)
-        boot_stats[i] = stat_fn(values[idx])
-    alpha = (1 - ci) / 2
-    return float(stat_fn(values)), float(np.percentile(boot_stats, alpha * 100)), \
-           float(np.percentile(boot_stats, (1 - alpha) * 100))
-
-
-def bootstrap_rsa(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    n_boot: int = 2000,
-    seed: int = 42,
-) -> Tuple[float, float, float]:
-    """Bootstrap CI for RSA Spearman r by resampling modules."""
-    rng = np.random.default_rng(seed)
-    n = module_emb.shape[0]
-    boot_rs = np.zeros(n_boot)
-    for i in range(n_boot):
-        idx = rng.choice(n, size=n, replace=True)
-        dist_p = pdist(module_emb[idx], metric="cosine")
-        dist_t = pdist(text_emb[idx], metric="cosine")
-        boot_rs[i], _ = spearmanr(dist_p, dist_t)
-    observed, _ = spearmanr(pdist(module_emb, "cosine"), pdist(text_emb, "cosine"))
-    return float(observed), float(np.percentile(boot_rs, 2.5)), \
-           float(np.percentile(boot_rs, 97.5))
-
-
-def get_biomod_sizes(module_json: str) -> Dict[str, int]:
-    """Get number of genes per module from merged_biomods.json."""
-    with open(module_json, "r", encoding="utf-8") as f:
-        modules = json.load(f)
-    return {name: len(genes) for name, genes in modules.items()}
-
-
-# ============================================================
-# CLI
-# ============================================================
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description=__doc__,
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    parser.add_argument("--biomod_embedding_path", "--module_embedding_path", type=str, required=True,
-                        help="Primary module_embeddings.npy (e.g. pretrain_subset)", dest='module_embedding_path')
-    parser.add_argument("--biomod_embedding_path2", "--module_embedding_path2", type=str, default=None,
-                        help="Second dataset module_embeddings.npy for cross-dataset comparison", dest='module_embedding_path2')
-    parser.add_argument("--biomod_json", "--module_json", type=str,
-                        default=str(REPO_ROOT / "model_dependencies" / "biomod_data_all" / "merged_biomods.json"), dest='module_json')
-    parser.add_argument("--descriptions_file", type=str, default=None)
-    parser.add_argument("--sentence_model", type=str,
-                        default="pritamdeka/S-PubMedBert-MS-MARCO")
-    parser.add_argument("--sentence_models_compare", type=str, default=None,
-                        help="Comma-separated additional models for comparison")
-    parser.add_argument("--device", type=str, default="cpu")
-    parser.add_argument("--n_permutations", type=int, default=1000)
-    parser.add_argument("--n_bootstrap", type=int, default=2000)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--output_dir", type=str, required=True)
-    return parser.parse_args()
-
-
-# ============================================================
-# Figure 1: RSA Overview (hexbin + null distribution)
-# ============================================================
-
-def fig1_rsa_overview(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    mantel_results: Dict,
-    output_dir: Path,
-    metric: str = "cosine",
-    n_bootstrap: int = 2000,
-    seed: int = 42,
-) -> Dict:
-    """RSA overview: hexbin density + Mantel null distribution."""
-    dist_module = pdist(module_emb, metric=metric)
-    dist_text = pdist(text_emb, metric=metric)
-
-    observed_r = mantel_results["rsa_spearman_observed"]
-    null_mean = mantel_results["rsa_null_mean"]
-    null_std = mantel_results["rsa_null_std"]
-    p_value = mantel_results["rsa_p_value"]
-    z_score = mantel_results["rsa_effect_size_z"]
-
-    # Bootstrap CI for RSA
-    rsa_obs, rsa_ci_lo, rsa_ci_hi = bootstrap_rsa(
-        module_emb, text_emb, n_boot=n_bootstrap, seed=seed
-    )
-
-    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
-
-    # Panel A: Hexbin density
-    ax = axes[0]
-    hb = ax.hexbin(dist_text, dist_module, gridsize=40, cmap="YlOrRd",
-                   mincnt=1, linewidths=0.2)
-    cb = plt.colorbar(hb, ax=ax, fraction=0.046, pad=0.02)
-    cb.set_label("Count", fontsize=9)
-
-    # Trend line
-    z = np.polyfit(dist_text, dist_module, 1)
-    p_line = np.poly1d(z)
-    x_range = np.linspace(dist_text.min(), dist_text.max(), 100)
-    ax.plot(x_range, p_line(x_range), "k-", linewidth=2, alpha=0.8)
-
-    ax.set_xlabel("Text pairwise distance (cosine)", fontsize=10)
-    ax.set_ylabel("Module embedding pairwise distance (cosine)", fontsize=10)
-    ax.set_title(f"A. RSA: r = {observed_r:.3f} [{rsa_ci_lo:.3f}, {rsa_ci_hi:.3f}]",
-                 fontsize=11, fontweight="bold")
-    ax.text(0.05, 0.95, f"n = {module_emb.shape[0]} modules\n"
-            f"pairs = {len(dist_module):,}",
-            transform=ax.transAxes, fontsize=8, va="top",
-            bbox=dict(boxstyle="round", facecolor="white", alpha=0.8))
-
-    # Panel B: Null distribution
-    ax = axes[1]
-    from scipy.stats import norm as norm_dist
-    x = np.linspace(null_mean - 4.5 * null_std,
-                    max(observed_r + 2 * null_std, null_mean + 4.5 * null_std), 300)
-    y = norm_dist.pdf(x, null_mean, null_std)
-    ax.fill_between(x, y, alpha=0.25, color="gray")
-    ax.plot(x, y, color="gray", linewidth=1.5, label="Null distribution")
-
-    # Shade p-value region
-    x_tail = x[x >= observed_r]
-    y_tail = norm_dist.pdf(x_tail, null_mean, null_std)
-    ax.fill_between(x_tail, y_tail, alpha=0.4, color="red", label=f"p = {p_value:.4f}")
-
-    ax.axvline(observed_r, color="red", linewidth=2.5, linestyle="-",
-               label=f"Observed r = {observed_r:.4f}")
-    ax.axvline(null_mean, color="gray", linewidth=1, linestyle="--", alpha=0.7)
-
-    ax.set_xlabel("Spearman r (distance matrix correlation)", fontsize=10)
-    ax.set_ylabel("Density", fontsize=10)
-    ax.set_title(f"B. Mantel Test: z = {z_score:.1f}σ, p = {p_value:.4f}",
-                 fontsize=11, fontweight="bold")
-    ax.legend(fontsize=9, loc="upper left")
-
-    plt.tight_layout()
-    savefig(fig, str(output_dir / "fig1_rsa_overview"))
-
-    return {
-        "rsa_r": rsa_obs,
-        "rsa_ci_95": [rsa_ci_lo, rsa_ci_hi],
-        "mantel_p": p_value,
-        "mantel_z": z_score,
-    }
-
-
-# ============================================================
-# Figure 2: Hallmarks Similarity Matrices with Dendrograms
-# ============================================================
-
-def _rank_transform_matrix(sim_matrix: np.ndarray) -> np.ndarray:
-    """Rank-transform off-diagonal elements to [0,1] for visualization."""
-    from scipy.stats import rankdata
-    n = sim_matrix.shape[0]
-    result = np.zeros_like(sim_matrix)
-    triu_idx = np.triu_indices(n, k=1)
-    vals = sim_matrix[triu_idx]
-    ranks = rankdata(vals) / len(vals)
-    result[triu_idx] = ranks
-    result[triu_idx[1], triu_idx[0]] = ranks
-    np.fill_diagonal(result, 1.0)
-    return result
-
-
-def fig2_hallmarks_heatmaps(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    output_dir: Path,
-) -> Dict:
-    """Hallmarks 50x50 similarity matrices with rank-transform and clustering."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    h_texts = [texts[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    h_processes = []
-    for n in h_names:
-        proc = parse_hallmark_process(n)
-        h_processes.append(proc if proc else "pathway")
-
-    module_sim = cosine_similarity(h_module)
-    text_sim = cosine_similarity(h_text)
-
-    linkage_mat = linkage(pdist(h_module, "cosine"), method="ward")
-    dendro = dendrogram(linkage_mat, no_plot=True)
-    order = dendro["leaves"]
-
-    module_sim_ord = module_sim[order][:, order]
-    text_sim_ord = text_sim[order][:, order]
-    texts_ord = [h_texts[i] for i in order]
-    procs_ord = [h_processes[i] for i in order]
-
-    # Rank-transform to spread narrow cosine similarity ranges
-    module_rank_ord = _rank_transform_matrix(module_sim_ord)
-    text_rank_ord = _rank_transform_matrix(text_sim_ord)
-    diff_rank = module_rank_ord - text_rank_ord
-
-    # Adaptive limit for difference panel
-    d_vals = diff_rank[np.triu_indices(n_h, k=1)]
-    vlim_d = float(np.percentile(np.abs(d_vals), 95))
-
-    # Figure
-    fig = plt.figure(figsize=(20, 8.5))
-    gs = fig.add_gridspec(2, 3, height_ratios=[0.08, 1], hspace=0.02, wspace=0.3)
-
-    # Dendrogram
-    ax_dend = fig.add_subplot(gs[0, :])
-    dendrogram(linkage_mat, ax=ax_dend, no_labels=True,
-               color_threshold=0, above_threshold_color="gray")
-    ax_dend.set_xticks([])
-    ax_dend.set_yticks([])
-    ax_dend.spines[:].set_visible(False)
-
-    for i, proc in enumerate(procs_ord):
-        color = PROCESS_COLORS.get(proc, "#999999")
-        ax_dend.axvspan(i * 10 + 5 - 4, i * 10 + 5 + 4, ymin=-0.2, ymax=-0.02,
-                        color=color, clip_on=False)
-
-    # Panel A: Module similarity (rank-transformed)
-    ax0 = fig.add_subplot(gs[1, 0])
-    im0 = ax0.imshow(module_rank_ord, cmap="YlOrRd", vmin=0, vmax=1,
-                     aspect="equal", interpolation="nearest")
-    ax0.set_title("A. Module Embedding Similarity\n(rank-transformed)",
-                  fontsize=10, fontweight="bold")
-    ax0.set_xticks(range(n_h))
-    ax0.set_xticklabels(texts_ord, rotation=90, fontsize=5.5, ha="center")
-    ax0.set_yticks(range(n_h))
-    ax0.set_yticklabels(texts_ord, fontsize=5.5)
-    cb0 = plt.colorbar(im0, ax=ax0, fraction=0.046, pad=0.02)
-    cb0.set_label("Rank percentile", fontsize=8)
-
-    # Panel B: Text similarity (rank-transformed)
-    ax1 = fig.add_subplot(gs[1, 1])
-    im1 = ax1.imshow(text_rank_ord, cmap="YlOrRd", vmin=0, vmax=1,
-                     aspect="equal", interpolation="nearest")
-    ax1.set_title("B. Text Embedding Similarity\n(rank-transformed)",
-                  fontsize=10, fontweight="bold")
-    ax1.set_xticks(range(n_h))
-    ax1.set_xticklabels(texts_ord, rotation=90, fontsize=5.5, ha="center")
-    ax1.set_yticks([])
-    cb1 = plt.colorbar(im1, ax=ax1, fraction=0.046, pad=0.02)
-    cb1.set_label("Rank percentile", fontsize=8)
-
-    # Panel C: Difference
-    ax2 = fig.add_subplot(gs[1, 2])
-    im2 = ax2.imshow(diff_rank, cmap="PiYG", vmin=-vlim_d, vmax=vlim_d,
-                     aspect="equal", interpolation="nearest")
-    ax2.set_title("C. Difference (Module − Text)\n(rank space)",
-                  fontsize=10, fontweight="bold")
-    ax2.set_xticks(range(n_h))
-    ax2.set_xticklabels(texts_ord, rotation=90, fontsize=5.5, ha="center")
-    ax2.set_yticks([])
-    cb2 = plt.colorbar(im2, ax=ax2, fraction=0.046, pad=0.02)
-    cb2.set_label("Rank difference", fontsize=8)
-
-    legend_elements = [LegendHandle(facecolor=PROCESS_COLORS[k], label=PROCESS_LABELS[k])
-                       for k in PROCESS_COLORS]
-    fig.legend(handles=legend_elements, loc="lower center", ncol=4,
-               fontsize=9, frameon=True, bbox_to_anchor=(0.5, -0.01))
-
-    plt.tight_layout(rect=[0, 0.03, 1, 1])
-    savefig(fig, str(output_dir / "fig2_hallmarks_heatmaps"))
-
-    return {
-        "module_sim_range": [float(module_sim.min()), float(module_sim.max())],
-        "text_sim_range": [float(text_sim.min()), float(text_sim.max())],
-        "clustering_order": [int(i) for i in order],
-    }
-
-
-# ============================================================
-# Figure 3: Per-Hallmark Alignment (lollipop chart)
-# ============================================================
-
-def fig3_per_hallmark_alignment(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    output_dir: Path,
-    n_bootstrap: int = 2000,
-    seed: int = 42,
-) -> Dict:
-    """Lollipop chart of per-hallmark alignment scores, colored by process."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    h_texts = [texts[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    h_processes = []
-    for n in h_names:
-        proc = parse_hallmark_process(n)
-        h_processes.append(proc if proc else "pathway")
-
-    # Per-module alignment: Spearman r of similarity profiles
-    module_sim = cosine_similarity(h_module)
-    text_sim = cosine_similarity(h_text)
-    per_module_corr = np.zeros(n_h)
-    for i in range(n_h):
-        mask = np.ones(n_h, dtype=bool)
-        mask[i] = False
-        r, _ = spearmanr(module_sim[i, mask], text_sim[i, mask])
-        per_module_corr[i] = r
-
-    # Sort by alignment
-    sort_idx = np.argsort(per_module_corr)  # ascending for bottom-to-top lollipop
-
-    # Bootstrap CI for mean alignment
-    mean_val, ci_lo, ci_hi = bootstrap_ci(per_module_corr, np.mean, n_bootstrap, seed=seed)
-    n_positive = int((per_module_corr > 0).sum())
-
-    # Figure
-    fig, ax = plt.subplots(figsize=(8, 12))
-
-    y_pos = np.arange(n_h)
-    for i, idx in enumerate(sort_idx):
-        color = PROCESS_COLORS.get(h_processes[idx], "#999999")
-        ax.hlines(y=i, xmin=0, xmax=per_module_corr[idx], color=color,
-                  linewidth=1.5, alpha=0.7)
-        ax.scatter(per_module_corr[idx], i, color=color, s=40, zorder=3)
-
-    ax.axvline(0, color="black", linewidth=0.8, linestyle="-")
-    ax.axvline(mean_val, color="black", linewidth=1.5, linestyle="--", alpha=0.6)
-
-    ax.set_yticks(y_pos)
-    ax.set_yticklabels([h_texts[i] for i in sort_idx], fontsize=7)
-    ax.set_xlabel("Alignment Score (Spearman r of similarity profiles)", fontsize=10)
-    ax.set_title("Per-Hallmark Alignment Score", fontsize=12, fontweight="bold")
-
-    # Summary text
-    summary = (f"Mean r = {mean_val:.3f} [{ci_lo:.3f}, {ci_hi:.3f}]\n"
-               f"Positive: {n_positive}/{n_h} ({100*n_positive/n_h:.0f}%)")
-    ax.text(0.98, 0.02, summary, transform=ax.transAxes, fontsize=9,
-            ha="right", va="bottom",
-            bbox=dict(boxstyle="round", facecolor="lightyellow", alpha=0.9))
-
-    # Legend
-    legend_elements = [Line2D([0], [0], marker="o", color="w", markerfacecolor=c,
-                              markersize=8, label=PROCESS_LABELS[k])
-                       for k, c in PROCESS_COLORS.items()]
-    ax.legend(handles=legend_elements, loc="upper left", fontsize=7,
-              frameon=True, ncol=2)
-
-    ax.set_xlim(min(per_module_corr) - 0.05, max(per_module_corr) + 0.05)
-    plt.tight_layout()
-    savefig(fig, str(output_dir / "fig3_per_hallmark_alignment"))
-
-    # Save CSV
-    df = pd.DataFrame({
-        "hallmark": h_texts,
-        "module_name": h_names,
-        "process_category": h_processes,
-        "alignment_r": per_module_corr,
-    }).sort_values("alignment_r", ascending=False)
-    df.to_csv(output_dir / "hallmarks_alignment_table.csv", index=False)
-
-    return {
-        "mean_alignment": mean_val,
-        "mean_alignment_ci_95": [ci_lo, ci_hi],
-        "n_positive": n_positive,
-        "n_total": n_h,
-        "per_module_scores": {h_texts[i]: float(per_module_corr[i]) for i in range(n_h)},
-    }
-
-
-# ============================================================
-# Figure 4: Alignment by Biological Process Category
-# ============================================================
-
-def fig4_category_alignment(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    output_dir: Path,
-) -> Dict:
-    """Violin/box plot of alignment scores grouped by MSigDB process category."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    h_processes = []
-    for n in h_names:
-        proc = parse_hallmark_process(n)
-        h_processes.append(proc if proc else "pathway")
-
-    # Per-module alignment
-    module_sim = cosine_similarity(h_module)
-    text_sim = cosine_similarity(h_text)
-    per_module_corr = np.zeros(n_h)
-    for i in range(n_h):
-        mask = np.ones(n_h, dtype=bool)
-        mask[i] = False
-        r, _ = spearmanr(module_sim[i, mask], text_sim[i, mask])
-        per_module_corr[i] = r
-
-    # Group by category
-    cat_scores = defaultdict(list)
-    for i in range(n_h):
-        cat_scores[h_processes[i]].append(per_module_corr[i])
-
-    # Sort categories by median alignment
-    cats_sorted = sorted(cat_scores.keys(),
-                         key=lambda c: np.median(cat_scores[c]), reverse=True)
-
-    # Kruskal-Wallis test
-    groups = [np.array(cat_scores[c]) for c in cats_sorted if len(cat_scores[c]) >= 2]
-    if len(groups) >= 2:
-        kw_stat, kw_p = kruskal(*groups)
-    else:
-        kw_stat, kw_p = 0.0, 1.0
-
-    # Figure
-    fig, ax = plt.subplots(figsize=(10, 6))
-
-    positions = range(len(cats_sorted))
-    data_list = [np.array(cat_scores[c]) for c in cats_sorted]
-    colors_list = [PROCESS_COLORS.get(c, "#999999") for c in cats_sorted]
-
-    # Box + strip plot
-    bp = ax.boxplot(data_list, positions=positions, widths=0.5,
-                    patch_artist=True, showfliers=False)
-    for module, color in zip(bp["boxes"], colors_list):
-        module.set_facecolor(color)
-        module.set_alpha(0.4)
-    for median in bp["medians"]:
-        median.set_color("black")
-        median.set_linewidth(1.5)
-
-    # Overlay individual points
-    for i, (data, color) in enumerate(zip(data_list, colors_list)):
-        jitter = np.random.default_rng(42).uniform(-0.15, 0.15, size=len(data))
-        ax.scatter(np.full(len(data), i) + jitter, data, color=color,
-                   s=30, alpha=0.7, edgecolors="white", linewidths=0.5, zorder=3)
-
-    ax.axhline(0, color="gray", linewidth=0.8, linestyle="--")
-    ax.set_xticks(positions)
-    ax.set_xticklabels([PROCESS_LABELS.get(c, c) for c in cats_sorted],
-                       rotation=30, ha="right", fontsize=9)
-    ax.set_ylabel("Alignment Score (Spearman r)", fontsize=10)
-    ax.set_title(f"Alignment by Biological Process Category\n"
-                 f"Kruskal-Wallis H = {kw_stat:.2f}, p = {kw_p:.4f}",
-                 fontsize=11, fontweight="bold")
-
-    # Annotate counts
-    for i, c in enumerate(cats_sorted):
-        n = len(cat_scores[c])
-        ax.text(i, ax.get_ylim()[1] * 0.95, f"n={n}", ha="center", fontsize=8)
-
-    plt.tight_layout()
-    savefig(fig, str(output_dir / "fig4_category_alignment"))
-
-    return {
-        "kruskal_wallis_H": float(kw_stat),
-        "kruskal_wallis_p": float(kw_p),
-        "category_medians": {c: float(np.median(cat_scores[c])) for c in cats_sorted},
-        "category_counts": {c: len(cat_scores[c]) for c in cats_sorted},
-    }
-
-
-# ============================================================
-# Figure 5: Granularity Effect (Reactome hierarchy)
-# ============================================================
-
-def fig5_granularity_effect(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    n_permutations: int = 1000,
-    n_bootstrap: int = 2000,
-    seed: int = 42,
-    output_dir: Path = None,
-) -> Dict:
-    """RSA at multiple granularity levels: fine → category → super-category."""
-    families = [parse_biomod_family(n) for n in names]
-    reactome_mask = np.array([f == "Reactome" for f in families])
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    reactome_idx = np.where(reactome_mask)[0]
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    results = {}
-
-    # Level 0: Hallmarks (50 modules, already shown to align)
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    r_hall, ci_lo_h, ci_hi_h = bootstrap_rsa(h_module, h_text, n_bootstrap, seed)
-    results["hallmarks_50"] = {"r": r_hall, "ci": [ci_lo_h, ci_hi_h], "n": len(hallmark_idx)}
-
-    # Level 1: Reactome fine-grained (491 pathways)
-    r_module = module_emb[reactome_idx]
-    r_text = text_emb[reactome_idx]
-    r_fine, ci_lo_f, ci_hi_f = bootstrap_rsa(r_module, r_text, n_bootstrap, seed)
-    results["reactome_fine_491"] = {"r": r_fine, "ci": [ci_lo_f, ci_hi_f], "n": len(reactome_idx)}
-
-    # Level 2: Reactome categories (~16)
-    r_names = [names[i] for i in reactome_idx]
-    r_texts = [texts[i] for i in reactome_idx]
-    cat_map = {}
-    for i, n in enumerate(r_names):
-        cat = categorize_reactome(n)
-        if cat not in cat_map:
-            cat_map[cat] = {"module_embs": [], "text_embs": []}
-        cat_map[cat]["module_embs"].append(r_module[i])
-        cat_map[cat]["text_embs"].append(r_text[i])
-
-    cats_with_data = {k: v for k, v in cat_map.items() if len(v["module_embs"]) >= 3}
-    if len(cats_with_data) >= 4:
-        cat_module_means = np.array([np.mean(v["module_embs"], axis=0)
-                                    for v in cats_with_data.values()])
-        cat_text_means = np.array([np.mean(v["text_embs"], axis=0)
-                                   for v in cats_with_data.values()])
-        r_cat, ci_lo_c, ci_hi_c = bootstrap_rsa(cat_module_means, cat_text_means,
-                                                  n_bootstrap, seed)
-        results["reactome_category"] = {
-            "r": r_cat, "ci": [ci_lo_c, ci_hi_c], "n": len(cats_with_data)
-        }
-    else:
-        results["reactome_category"] = {"r": 0.0, "ci": [0.0, 0.0], "n": 0}
-
-    # Level 3: Super-categories (5-6)
-    super_map = {}
-    for super_name, sub_cats in REACTOME_SUPER_CATEGORIES.items():
-        embs_p, embs_t = [], []
-        for sub in sub_cats:
-            if sub in cats_with_data:
-                embs_p.extend(cats_with_data[sub]["module_embs"])
-                embs_t.extend(cats_with_data[sub]["text_embs"])
-        if len(embs_p) >= 3:
-            super_map[super_name] = {
-                "module_mean": np.mean(embs_p, axis=0),
-                "text_mean": np.mean(embs_t, axis=0),
-            }
-
-    if len(super_map) >= 4:
-        sup_module = np.array([v["module_mean"] for v in super_map.values()])
-        sup_text = np.array([v["text_mean"] for v in super_map.values()])
-        r_sup, ci_lo_s, ci_hi_s = bootstrap_rsa(sup_module, sup_text, n_bootstrap, seed)
-        results["reactome_super_category"] = {
-            "r": r_sup, "ci": [ci_lo_s, ci_hi_s], "n": len(super_map)
-        }
-    else:
-        results["reactome_super_category"] = {"r": 0.0, "ci": [0.0, 0.0], "n": 0}
-
-    # Figure
-    fig, ax = plt.subplots(figsize=(8, 5))
-
-    levels = ["Reactome\nFine (491)", "Reactome\nCategory (~16)",
-              "Reactome\nSuper-Cat (~6)", "Hallmarks\n(50)"]
-    keys = ["reactome_fine_491", "reactome_category",
-            "reactome_super_category", "hallmarks_50"]
-    rs = [results[k]["r"] for k in keys]
-    cis_lo = [results[k]["ci"][0] for k in keys]
-    cis_hi = [results[k]["ci"][1] for k in keys]
-    ns = [results[k]["n"] for k in keys]
-
-    colors = ["#90CAF9", "#42A5F5", "#1565C0", "#E91E63"]
-    x = np.arange(len(levels))
-    bars = ax.bar(x, rs, color=colors, edgecolor="black", linewidth=0.8, width=0.6)
-
-    # Error bars (CI)
-    for i in range(len(levels)):
-        ax.plot([i, i], [cis_lo[i], cis_hi[i]], color="black", linewidth=2)
-        ax.plot([i - 0.1, i + 0.1], [cis_lo[i], cis_lo[i]], color="black", linewidth=1.5)
-        ax.plot([i - 0.1, i + 0.1], [cis_hi[i], cis_hi[i]], color="black", linewidth=1.5)
-
-    ax.axhline(0, color="gray", linewidth=0.8, linestyle="--")
-    ax.set_xticks(x)
-    ax.set_xticklabels(levels, fontsize=9)
-    ax.set_ylabel("RSA Spearman r", fontsize=10)
-    ax.set_title("Alignment Improves at Coarser Granularity",
-                 fontsize=12, fontweight="bold")
-
-    # Annotate n
-    for i, n in enumerate(ns):
-        ax.text(i, max(rs[i], cis_hi[i]) + 0.01, f"n={n}", ha="center", fontsize=8)
-
-    plt.tight_layout()
-    savefig(fig, str(output_dir / "fig5_granularity_effect"))
-
-    return results
-
-
-# ============================================================
-# Figure 6: Multi-Model / Cross-Dataset Comparison
-# ============================================================
-
-def fig6_model_comparison(
-    module_emb_primary: np.ndarray,
-    module_emb_secondary: Optional[np.ndarray],
-    text_emb_primary: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    primary_model_name: str,
-    compare_model_names: List[str],
-    compare_text_embs: List[np.ndarray],
-    n_permutations: int = 1000,
-    n_bootstrap: int = 2000,
-    seed: int = 42,
-    output_dir: Path = None,
-) -> Dict:
-    """Grouped bar chart comparing RSA across models and datasets."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb_primary[hallmark_idx]
-
-    # Compute RSA for ALL models including primary
-    all_models = [primary_model_name] + compare_model_names
-    all_text_embs = [text_emb_primary] + compare_text_embs
-
-    model_results = {}
-    for model_name, text_emb_full in zip(all_models, all_text_embs):
-        h_text = text_emb_full[hallmark_idx]
-        r, ci_lo, ci_hi = bootstrap_rsa(h_module, h_text, n_bootstrap, seed)
-        mantel_res = mantel_test(h_module, h_text, n_perm=n_permutations)
-        p = mantel_res["rsa_p_value"]
-        model_results[model_name] = {
-            "dataset": "primary",
-            "r": r, "ci": [ci_lo, ci_hi], "p": p
-        }
-
-    # Cross-dataset (secondary)
-    dataset_results = {}
-    if module_emb_secondary is not None:
-        h_module2 = module_emb_secondary[hallmark_idx]
-        for model_name, text_emb_full in zip(all_models, all_text_embs):
-            h_text = text_emb_full[hallmark_idx]
-            r, ci_lo, ci_hi = bootstrap_rsa(h_module2, h_text, n_bootstrap, seed)
-            mantel_res = mantel_test(h_module2, h_text, n_perm=n_permutations)
-            p = mantel_res["rsa_p_value"]
-            dataset_results[model_name] = {
-                "dataset": "secondary",
-                "r": r, "ci": [ci_lo, ci_hi], "p": p
-            }
-
-    # Figure
-    has_secondary = len(dataset_results) > 0
-    n_models = len(model_results)
-
-    if n_models == 0:
-        logger.warning("No models to compare, skipping fig6")
-        return {}
-
-    fig, ax = plt.subplots(figsize=(max(8, n_models * 2.5), 5))
-
-    model_names_short = [m.split("/")[-1][:25] for m in model_results.keys()]
-    x = np.arange(len(model_names_short))
-    width = 0.35 if has_secondary else 0.5
-
-    # Primary dataset bars
-    rs_primary = [model_results[m]["r"] for m in model_results]
-    cis_primary = [model_results[m]["ci"] for m in model_results]
-    yerr_lo = [r - ci[0] for r, ci in zip(rs_primary, cis_primary)]
-    yerr_hi = [ci[1] - r for r, ci in zip(rs_primary, cis_primary)]
-
-    # Color: highlight primary model differently
-    colors_primary = ["#1E88E5"] + ["#90CAF9"] * (n_models - 1)
-
-    offset = -width / 2 if has_secondary else 0
-    bars1 = ax.bar(x + offset, rs_primary, width, label="Dataset 1 (pretrain)",
-                   color=colors_primary, edgecolor="black", linewidth=0.8)
-    ax.errorbar(x + offset, rs_primary, yerr=[yerr_lo, yerr_hi],
-                fmt="none", color="black", capsize=4)
-
-    if has_secondary:
-        rs_sec = [dataset_results.get(m, {}).get("r", 0) for m in model_results]
-        cis_sec = [dataset_results.get(m, {}).get("ci", [0, 0]) for m in model_results]
-        yerr_lo2 = [r - ci[0] for r, ci in zip(rs_sec, cis_sec)]
-        yerr_hi2 = [ci[1] - r for r, ci in zip(rs_sec, cis_sec)]
-        colors_sec = ["#E53935"] + ["#EF9A9A"] * (n_models - 1)
-        bars2 = ax.bar(x + width / 2, rs_sec, width, label="Dataset 2 (pancreas)",
-                       color=colors_sec, edgecolor="black", linewidth=0.8)
-        ax.errorbar(x + width / 2, rs_sec, yerr=[yerr_lo2, yerr_hi2],
-                    fmt="none", color="black", capsize=4)
-
-    ax.axhline(0, color="gray", linewidth=0.8, linestyle="--")
-    ax.set_xticks(x)
-    ax.set_xticklabels(model_names_short, rotation=15, ha="right", fontsize=9)
-    ax.set_ylabel("RSA Spearman r (Hallmarks)", fontsize=10)
-    ax.set_title("Alignment Across Sentence Encoders",
-                 fontsize=12, fontweight="bold")
-    ax.legend(fontsize=9)
-
-    # Annotate p-values
-    for i, m in enumerate(model_results):
-        p = model_results[m]["p"]
-        sig = "***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else "ns"
-        y_pos = max(rs_primary[i], cis_primary[i][1]) + 0.02
-        ax.text(i + offset, y_pos, sig, ha="center", fontsize=9, fontweight="bold")
-
-    plt.tight_layout()
-    savefig(fig, str(output_dir / "fig6_model_comparison"))
-
-    return {"model_results": model_results, "dataset_results": dataset_results}
-
-
-# ============================================================
-# Deep Analysis: Alignment vs Module Size
-# ============================================================
-
-def analysis_alignment_vs_biomod_size(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    module_json: str,
-    output_dir: Path,
-) -> Dict:
-    """Scatter: n_genes vs alignment score. Tests if larger modules align better."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    h_texts = [texts[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    # Per-module alignment
-    module_sim = cosine_similarity(h_module)
-    text_sim = cosine_similarity(h_text)
-    per_module_corr = np.zeros(n_h)
-    for i in range(n_h):
-        mask = np.ones(n_h, dtype=bool)
-        mask[i] = False
-        r, _ = spearmanr(module_sim[i, mask], text_sim[i, mask])
-        per_module_corr[i] = r
-
-    # Module sizes
-    module_sizes = get_biomod_sizes(module_json)
-    h_sizes = np.array([module_sizes.get(n, 0) for n in h_names])
-
-    # Correlation
-    r_size, p_size = spearmanr(h_sizes, per_module_corr)
-
-    # Figure
-    fig, ax = plt.subplots(figsize=(7, 5))
-    h_processes = [parse_hallmark_process(n) or "pathway" for n in h_names]
-    colors = [PROCESS_COLORS.get(p, "#999999") for p in h_processes]
-
-    ax.scatter(h_sizes, per_module_corr, c=colors, s=50, alpha=0.7,
-               edgecolors="white", linewidths=0.5)
-
-    # Trend line
-    z = np.polyfit(h_sizes, per_module_corr, 1)
-    p_line = np.poly1d(z)
-    x_range = np.linspace(h_sizes.min(), h_sizes.max(), 100)
-    ax.plot(x_range, p_line(x_range), "k--", linewidth=1.5, alpha=0.6)
-
-    ax.axhline(0, color="gray", linewidth=0.8, linestyle=":")
-    ax.set_xlabel("Number of Genes in Module", fontsize=10)
-    ax.set_ylabel("Alignment Score (Spearman r)", fontsize=10)
-    ax.set_title(f"Alignment vs Module Size\n"
-                 f"Spearman ρ = {r_size:.3f}, p = {p_size:.4f}",
-                 fontsize=11, fontweight="bold")
-
-    # Label outliers
-    for i in range(n_h):
-        if abs(per_module_corr[i]) > 0.35 or h_sizes[i] > 250:
-            ax.annotate(h_texts[i][:15], (h_sizes[i], per_module_corr[i]),
-                        fontsize=6, alpha=0.8, xytext=(5, 3),
-                        textcoords="offset points")
-
-    plt.tight_layout()
-    savefig(fig, str(output_dir / "fig_alignment_vs_size"))
-
-    return {"spearman_r": float(r_size), "spearman_p": float(p_size)}
-
-
-# ============================================================
-# Deep Analysis: Negative Result Deep-Dive
-# ============================================================
-
-def analysis_negative_results(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    texts: List[str],
-    output_dir: Path,
-    k: int = 5,
-    n_show: int = 5,
-) -> Dict:
-    """For poorly-aligned modules, show neighbor divergence."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    h_texts = [texts[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    # Per-module alignment
-    module_sim = cosine_similarity(h_module)
-    text_sim = cosine_similarity(h_text)
-    per_module_corr = np.zeros(n_h)
-    for i in range(n_h):
-        mask = np.ones(n_h, dtype=bool)
-        mask[i] = False
-        r, _ = spearmanr(module_sim[i, mask], text_sim[i, mask])
-        per_module_corr[i] = r
-
-    # Bottom n_show and top n_show
-    bottom_idx = np.argsort(per_module_corr)[:n_show]
-    top_idx = np.argsort(per_module_corr)[-n_show:][::-1]
-
-    rows = []
-    for idx_set, label in [(bottom_idx, "bottom"), (top_idx, "top")]:
-        for i in idx_set:
-            # kNN in module space
-            module_nn = np.argsort(-module_sim[i])
-            module_nn = [j for j in module_nn if j != i][:k]
-            # kNN in text space
-            text_nn = np.argsort(-text_sim[i])
-            text_nn = [j for j in text_nn if j != i][:k]
-            overlap = set(module_nn) & set(text_nn)
-
-            rows.append({
-                "rank": label,
-                "hallmark": h_texts[i],
-                "alignment_r": float(per_module_corr[i]),
-                "module_neighbors": ", ".join([h_texts[j] for j in module_nn]),
-                "text_neighbors": ", ".join([h_texts[j] for j in text_nn]),
-                "knn_overlap": len(overlap),
-            })
-
-    df = pd.DataFrame(rows)
-    df.to_csv(output_dir / "negative_analysis.csv", index=False)
-    logger.info("Saved negative_analysis.csv (%d rows)", len(df))
-
-    return {
-        "bottom_5": [{
-            "hallmark": r["hallmark"],
-            "r": r["alignment_r"],
-            "module_nn": r["module_neighbors"],
-            "text_nn": r["text_neighbors"],
-        } for _, r in df[df["rank"] == "bottom"].iterrows()],
-        "top_5": [{
-            "hallmark": r["hallmark"],
-            "r": r["alignment_r"],
-        } for _, r in df[df["rank"] == "top"].iterrows()],
-    }
-
-
-# ============================================================
-# Deep Analysis: kNN Overlap with Bootstrap CI
-# ============================================================
-
-def analysis_knn_overlap(
-    module_emb: np.ndarray,
-    text_emb: np.ndarray,
-    names: List[str],
-    n_permutations: int = 1000,
-    n_bootstrap: int = 2000,
-    k: int = 5,
-    seed: int = 42,
-    output_dir: Path = None,
-) -> Dict:
-    """kNN overlap with bootstrap CI and per-category breakdown."""
-    families = [parse_biomod_family(n) for n in names]
-    hallmark_mask = np.array([f == "Hallmarks" for f in families])
-    hallmark_idx = np.where(hallmark_mask)[0]
-
-    h_module = module_emb[hallmark_idx]
-    h_text = text_emb[hallmark_idx]
-    h_names = [names[i] for i in hallmark_idx]
-    n_h = len(hallmark_idx)
-
-    # Compute kNN overlap
-    module_sim = cosine_similarity(h_module)
-    text_sim = cosine_similarity(h_text)
-
-    overlaps = np.zeros(n_h)
-    for i in range(n_h):
-        module_nn = set(np.argsort(-module_sim[i])[1:k + 1])
-        text_nn = set(np.argsort(-text_sim[i])[1:k + 1])
-        overlaps[i] = len(module_nn & text_nn) / k
-
-    # Bootstrap CI
-    mean_obs, ci_lo, ci_hi = bootstrap_ci(overlaps, np.mean, n_bootstrap, seed=seed)
-
-    # Null distribution (permutation)
-    rng = np.random.default_rng(seed)
-    null_means = np.zeros(n_permutations)
-    for p in range(n_permutations):
-        perm = rng.permutation(n_h)
-        text_sim_perm = text_sim[perm][:, perm]
-        ov_perm = np.zeros(n_h)
-        for i in range(n_h):
-            module_nn = set(np.argsort(-module_sim[i])[1:k + 1])
-            text_nn = set(np.argsort(-text_sim_perm[i])[1:k + 1])
-            ov_perm[i] = len(module_nn & text_nn) / k
-        null_means[p] = ov_perm.mean()
-
-    p_value = float((null_means >= mean_obs).sum() / n_permutations)
-    null_mean = float(null_means.mean())
-    null_std = float(null_means.std())
-    z_score = (mean_obs - null_mean) / null_std if null_std > 0 else 0.0
-
-    # Per-category breakdown
-    h_processes = [parse_hallmark_process(n) or "pathway" for n in h_names]
-    cat_overlaps = defaultdict(list)
-    for i in range(n_h):
-        cat_overlaps[h_processes[i]].append(overlaps[i])
-
-    return {
-        "knn_overlap_mean": mean_obs,
-        "knn_overlap_ci_95": [ci_lo, ci_hi],
-        "null_mean": null_mean,
-        "null_std": null_std,
-        "p_value": p_value,
-        "z_score": z_score,
-        "k": k,
-        "category_means": {c: float(np.mean(v)) for c, v in cat_overlaps.items()},
-    }
-
-
-# ============================================================
-# Main
-# ============================================================
-
-def main():
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Setup logging
-    log_file = output_dir / "run_publication.log"
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s [%(levelname)s] %(message)s",
-        handlers=[
-            logging.FileHandler(log_file, mode="w", encoding="utf-8"),
-            logging.StreamHandler(),
-        ],
-    )
-    logger.info("Publication analysis starting")
-    logger.info("Args: %s", vars(args))
-
-    # ---- Load data ----
-    module_names = load_biomod_names(args.module_json)
-    all_module_emb = np.load(args.module_embedding_path)
-    logger.info("Loaded module embeddings: %s from %s",
-                all_module_emb.shape, args.module_embedding_path)
-
-    # Filter to semantic modules (Hallmarks + Reactome)
-    families = [parse_biomod_family(n) for n in module_names]
-    semantic_mask = np.array([f in ("Hallmarks", "Reactome") for f in families])
-    semantic_idx = np.where(semantic_mask)[0]
-    module_emb = all_module_emb[semantic_idx]
-    names = [module_names[i] for i in semantic_idx]
-    logger.info("Semantic modules: %d (Hallmarks + Reactome)", len(names))
-
-    # ---- Build text descriptions and encode ----
-    desc_file = getattr(args, "descriptions_file", None)
-    _, _, texts = build_text_descriptions(names, descriptions_file=desc_file)
-    logger.info("Encoding %d texts with primary model: %s", len(texts), args.sentence_model)
-    text_emb = encode_texts(texts, args.sentence_model, device=args.device)
-    logger.info("Text embeddings: %s", text_emb.shape)
-
-    # ---- Core metrics (RSA + Mantel) ----
-    logger.info("Computing RSA + Mantel test (n_perm=%d)...", args.n_permutations)
-    mantel_results = mantel_test(module_emb, text_emb, n_perm=args.n_permutations)
-    logger.info("RSA r=%.4f, p=%.4f, z=%.1f",
-                mantel_results["rsa_spearman_observed"],
-                mantel_results["rsa_p_value"],
-                mantel_results["rsa_effect_size_z"])
-
-    # ---- Figure 1: RSA Overview ----
-    logger.info("Generating Figure 1: RSA Overview...")
-    metrics_fig1 = fig1_rsa_overview(
-        module_emb, text_emb, mantel_results, output_dir,
-        n_bootstrap=args.n_bootstrap, seed=args.seed,
-    )
-
-    # ---- Figure 2: Hallmarks Heatmaps ----
-    logger.info("Generating Figure 2: Hallmarks Heatmaps...")
-    metrics_fig2 = fig2_hallmarks_heatmaps(
-        module_emb, text_emb, names, texts, output_dir,
-    )
-
-    # ---- Figure 3: Per-Hallmark Alignment ----
-    logger.info("Generating Figure 3: Per-Hallmark Alignment...")
-    metrics_fig3 = fig3_per_hallmark_alignment(
-        module_emb, text_emb, names, texts, output_dir,
-        n_bootstrap=args.n_bootstrap, seed=args.seed,
-    )
-
-    # ---- Figure 4: Category Alignment ----
-    logger.info("Generating Figure 4: Category Alignment...")
-    metrics_fig4 = fig4_category_alignment(
-        module_emb, text_emb, names, texts, output_dir,
-    )
-
-    # ---- Figure 5: Granularity Effect ----
-    logger.info("Generating Figure 5: Granularity Effect...")
-    metrics_fig5 = fig5_granularity_effect(
-        module_emb, text_emb, names, texts,
-        n_permutations=args.n_permutations,
-        n_bootstrap=args.n_bootstrap,
-        seed=args.seed,
-        output_dir=output_dir,
-    )
-
-    # ---- Figure 6: Multi-Model Comparison (optional) ----
-    metrics_fig6 = {}
-    compare_models = []
-    compare_embs = []
-    if args.sentence_models_compare:
-        compare_models = [m.strip() for m in args.sentence_models_compare.split(",")
-                          if m.strip()]
-        logger.info("Loading %d comparison models...", len(compare_models))
-        for model_name in compare_models:
-            logger.info("  Encoding with: %s", model_name)
-            try:
-                emb = encode_texts(texts, model_name, device=args.device)
-                compare_embs.append(emb)
-            except Exception as e:
-                logger.warning("  Failed to load %s: %s", model_name, e)
-                compare_models.remove(model_name)
-
-    module_emb2 = None
-    if args.module_embedding_path2:
-        all_module_emb2 = np.load(args.module_embedding_path2)
-        module_emb2 = all_module_emb2[semantic_idx]
-        logger.info("Loaded secondary dataset: %s", all_module_emb2.shape)
-
-    if compare_models:
-        logger.info("Generating Figure 6: Model Comparison...")
-        metrics_fig6 = fig6_model_comparison(
-            module_emb, module_emb2, text_emb, names, texts,
-            primary_model_name=args.sentence_model.split("/")[-1],
-            compare_model_names=compare_models,
-            compare_text_embs=compare_embs,
-            n_permutations=args.n_permutations,
-            n_bootstrap=args.n_bootstrap,
-            seed=args.seed,
-            output_dir=output_dir,
-        )
-
-    # ---- Deep Analysis ----
-    logger.info("Running deep analysis: alignment vs module size...")
-    metrics_size = analysis_alignment_vs_biomod_size(
-        module_emb, text_emb, names, texts, args.module_json, output_dir,
-    )
-
-    logger.info("Running deep analysis: negative results...")
-    metrics_neg = analysis_negative_results(
-        module_emb, text_emb, names, texts, output_dir,
-    )
-
-    logger.info("Running deep analysis: kNN overlap with bootstrap...")
-    metrics_knn = analysis_knn_overlap(
-        module_emb, text_emb, names,
-        n_permutations=args.n_permutations,
-        n_bootstrap=args.n_bootstrap,
-        seed=args.seed,
-        output_dir=output_dir,
-    )
-    logger.info("kNN overlap: %.3f [%.3f, %.3f], null=%.3f, p=%.4f",
-                metrics_knn["knn_overlap_mean"],
-                metrics_knn["knn_overlap_ci_95"][0],
-                metrics_knn["knn_overlap_ci_95"][1],
-                metrics_knn["null_mean"],
-                metrics_knn["p_value"])
-
-    # ---- Cross-dataset RSA (if secondary available) ----
-    metrics_cross = {}
-    if module_emb2 is not None:
-        logger.info("Cross-dataset RSA (secondary dataset)...")
-        mantel2 = mantel_test(module_emb2, text_emb, n_perm=args.n_permutations)
-        r2, ci2_lo, ci2_hi = bootstrap_rsa(
-            module_emb2, text_emb, args.n_bootstrap, args.seed
-        )
-        metrics_cross = {
-            "secondary_rsa_r": r2,
-            "secondary_rsa_ci_95": [ci2_lo, ci2_hi],
-            "secondary_mantel_p": mantel2["rsa_p_value"],
-            "secondary_mantel_z": mantel2["rsa_effect_size_z"],
-        }
-        logger.info("Secondary RSA: r=%.4f [%.4f, %.4f], p=%.4f",
-                    r2, ci2_lo, ci2_hi, mantel2["rsa_p_value"])
-
-    # ---- Aggregate and save metrics ----
-    all_metrics = {
-        "primary_dataset": args.module_embedding_path,
-        "secondary_dataset": args.module_embedding_path2,
-        "sentence_model": args.sentence_model,
-        "n_modules": len(names),
-        "n_permutations": args.n_permutations,
-        "n_bootstrap": args.n_bootstrap,
-        "fig1_rsa_overview": metrics_fig1,
-        "fig2_heatmaps": metrics_fig2,
-        "fig3_per_hallmark": metrics_fig3,
-        "fig4_category": metrics_fig4,
-        "fig5_granularity": metrics_fig5,
-        "fig6_model_comparison": metrics_fig6,
-        "alignment_vs_size": metrics_size,
-        "negative_analysis": metrics_neg,
-        "knn_overlap": metrics_knn,
-        "cross_dataset": metrics_cross,
-    }
-
-    metrics_path = output_dir / "publication_metrics.json"
-    with open(metrics_path, "w", encoding="utf-8") as f:
-        json.dump(all_metrics, f, indent=2, ensure_ascii=False, default=str)
-    logger.info("Saved: %s", metrics_path)
-    logger.info("Publication analysis complete. Output: %s", output_dir)
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/biomod_text_alignment/fetch_reactome_descriptions.py

@@ -1,97 +0,0 @@
-"""Fetch Reactome pathway descriptions via the Content Service REST API.
-
-Produces a JSON file mapping module_name -> summaryText that can be passed
-to biomod_text_alignment_extended.py via --descriptions_file.
-
-Usage (run from the BioMod repository root):
-    python BioMod/interpretability/biomod_text_alignment/fetch_reactome_descriptions.py \
-        --biomod_json model_dependencies/biomod_data_all/merged_biomods.json \
-        --output reactome_descriptions.json
-
-Requires internet access to https://reactome.org/ContentService/
-"""
-import argparse
-import json
-import sys
-import time
-from pathlib import Path
-
-try:
-    import requests
-except ImportError:
-    print("ERROR: 'requests' package required. Install: pip install requests")
-    sys.exit(1)
-
-HERE = Path(__file__).resolve().parent
-INTERP_DIR = HERE.parent
-BIOMOD_SOURCE = INTERP_DIR.parent
-for p in (str(INTERP_DIR), str(BIOMOD_SOURCE)):
-    if p not in sys.path:
-        sys.path.insert(0, p)
-
-from interpretability.utils import load_biomod_names, parse_biomod_family
-
-REACTOME_API = "https://reactome.org/ContentService/data/query"
-
-
-def fetch_summary(rhs_id: str, retries: int = 3) -> str:
-    """Fetch summaryText for a Reactome pathway."""
-    url = f"{REACTOME_API}/{rhs_id}"
-    for attempt in range(retries):
-        try:
-            resp = requests.get(url, timeout=10)
-            if resp.status_code == 200:
-                data = resp.json()
-                return data.get("summation", [{}])[0].get("text", "")
-            elif resp.status_code == 404:
-                return ""
-            elif resp.status_code == 429:
-                time.sleep(2 ** attempt)
-            else:
-                return ""
-        except (requests.RequestException, json.JSONDecodeError):
-            time.sleep(1)
-    return ""
-
-
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--biomod_json", "--module_json", type=str, required=True, dest='module_json')
-    parser.add_argument("--output", type=str, default="reactome_descriptions.json")
-    parser.add_argument("--delay", type=float, default=0.1,
-                        help="Delay between API calls (seconds)")
-    args = parser.parse_args()
-
-    module_names = load_biomod_names(args.module_json)
-    reactome_modules = [n for n in module_names if parse_biomod_family(n) == "Reactome"]
-    print(f"Reactome modules to fetch: {len(reactome_modules)}")
-
-    descriptions = {}
-    for i, pname in enumerate(reactome_modules):
-        raw = pname.split("__", 1)[1]
-        parts = raw.rsplit(" R-HSA-", 1)
-        if len(parts) != 2:
-            continue
-        rhs_id = "R-HSA-" + parts[1]
-
-        summary = fetch_summary(rhs_id)
-        if summary:
-            # Strip HTML tags
-            import re
-            summary = re.sub(r"<[^>]+>", "", summary).strip()
-            descriptions[pname] = summary
-
-        if (i + 1) % 50 == 0:
-            print(f"  Fetched {i+1}/{len(reactome_modules)}, "
-                  f"got descriptions for {len(descriptions)}")
-        time.sleep(args.delay)
-
-    print(f"\nDone. Got descriptions for {len(descriptions)}/{len(reactome_modules)} pathways.")
-
-    with open(args.output, "w", encoding="utf-8") as f:
-        json.dump(descriptions, f, indent=2, ensure_ascii=False)
-    print(f"Saved: {args.output}")
-
-
-if __name__ == "__main__":
-    main()

BioMod/interpretability/utils.py

@@ -1,200 +0,0 @@
-"""Shared helpers for module interpretability experiments.
-
-Provides:
-    * module family / hallmark-process parsing from merged_biomods.json keys
-    * a generic in-memory dataset that yields raw [G_input] expression rows,
-      compatible with DataPreprocessor.reindex_tensor_to_model_space
-"""
-from __future__ import annotations
-
-import json
-import logging
-import os
-from pathlib import Path
-from typing import Dict, List, Optional, Sequence
-
-import numpy as np
-import torch
-from scipy import sparse
-from torch.utils.data import Dataset
-
-logger = logging.getLogger(__name__)
-
-
-# ---------- 1. Module family parsing ----------
-
-MODULE_FAMILY_MAP: Dict[str, str] = {
-    "module1_hallmarks": "Hallmarks",
-    "module2_reactome_2022_filtered": "Reactome",
-    "module3_hotspot_260k": "Hotspot",
-    "module3_processed": "Processed",
-}
-
-# MSigDB hallmark official process categories. Reference:
-# https://www.gsea-msigdb.org/gsea/msigdb/human/genesets.jsp, collection=H
-HALLMARK_PROCESS: Dict[str, str] = {
-    "ADIPOGENESIS": "metabolic",
-    "ALLOGRAFT_REJECTION": "immune",
-    "ANDROGEN_RESPONSE": "signaling",
-    "ANGIOGENESIS": "development",
-    "APICAL_JUNCTION": "cellular_component",
-    "APICAL_SURFACE": "cellular_component",
-    "APOPTOSIS": "pathway",
-    "BILE_ACID_METABOLISM": "metabolic",
-    "CHOLESTEROL_HOMEOSTASIS": "metabolic",
-    "COAGULATION": "immune",
-    "COMPLEMENT": "immune",
-    "DNA_REPAIR": "dna_damage",
-    "E2F_TARGETS": "proliferation",
-    "EPITHELIAL_MESENCHYMAL_TRANSITION": "development",
-    "ESTROGEN_RESPONSE_EARLY": "signaling",
-    "ESTROGEN_RESPONSE_LATE": "signaling",
-    "FATTY_ACID_METABOLISM": "metabolic",
-    "G2M_CHECKPOINT": "proliferation",
-    "GLYCOLYSIS": "metabolic",
-    "HEDGEHOG_SIGNALING": "signaling",
-    "HEME_METABOLISM": "metabolic",
-    "HYPOXIA": "pathway",
-    "IL2_STAT5_SIGNALING": "immune",
-    "IL6_JAK_STAT3_SIGNALING": "immune",
-    "INFLAMMATORY_RESPONSE": "immune",
-    "INTERFERON_ALPHA_RESPONSE": "immune",
-    "INTERFERON_GAMMA_RESPONSE": "immune",
-    "KRAS_SIGNALING_DN": "signaling",
-    "KRAS_SIGNALING_UP": "signaling",
-    "MITOTIC_SPINDLE": "proliferation",
-    "MTORC1_SIGNALING": "signaling",
-    "MYC_TARGETS_V1": "proliferation",
-    "MYC_TARGETS_V2": "proliferation",
-    "MYOGENESIS": "development",
-    "NOTCH_SIGNALING": "signaling",
-    "OXIDATIVE_PHOSPHORYLATION": "metabolic",
-    "P53_PATHWAY": "proliferation",
-    "PANCREAS_BETA_CELLS": "development",
-    "PEROXISOME": "cellular_component",
-    "PI3K_AKT_MTOR_SIGNALING": "signaling",
-    "PROTEIN_SECRETION": "pathway",
-    "REACTIVE_OXYGEN_SPECIES_PATHWAY": "pathway",
-    "SPERMATOGENESIS": "development",
-    "TGF_BETA_SIGNALING": "signaling",
-    "TNFA_SIGNALING_VIA_NFKB": "immune",
-    "UNFOLDED_PROTEIN_RESPONSE": "pathway",
-    "UV_RESPONSE_DN": "dna_damage",
-    "UV_RESPONSE_UP": "dna_damage",
-    "WNT_BETA_CATENIN_SIGNALING": "signaling",
-    "XENOBIOTIC_METABOLISM": "metabolic",
-}
-
-
-def parse_biomod_family(module_name: str) -> str:
-    head = module_name.split("__", 1)[0]
-    return MODULE_FAMILY_MAP.get(head, "Other")
-
-
-def parse_hallmark_process(module_name: str) -> Optional[str]:
-    if not module_name.startswith("module1_hallmarks__"):
-        return None
-    short = module_name.split("__", 1)[1].replace("HALLMARK_", "", 1)
-    return HALLMARK_PROCESS.get(short)
-
-
-def build_biomod_metadata(module_names: Sequence[str]):
-    rows = []
-    for name in module_names:
-        rows.append({
-            "module_name": name,
-            "family": parse_biomod_family(name),
-            "hallmark_process": parse_hallmark_process(name),
-        })
-    return rows
-
-
-def load_biomod_names(module_json_path: str) -> List[str]:
-    with open(module_json_path, "r", encoding="utf-8") as handle:
-        return list(json.load(handle).keys())
-
-
-# ---------- 2. Generic in-memory dataset ----------
-
-class MatrixRowDataset(Dataset):
-    """Yield dense float32 rows from a (sparse or dense) [N, G] matrix.
-
-    Used for inference-only flows where DataPreprocessor.reindex_tensor_to_model_space
-    will project [B, G_input] -> [B, G_model].
-    """
-
-    def __init__(self, matrix):
-        if sparse.issparse(matrix):
-            self.matrix = matrix.tocsr()
-            self._sparse = True
-        else:
-            self.matrix = np.asarray(matrix, dtype=np.float32)
-            self._sparse = False
-        self.n = self.matrix.shape[0]
-
-    def __len__(self) -> int:
-        return self.n
-
-    def __getitem__(self, idx: int) -> torch.Tensor:
-        row = self.matrix[idx]
-        if self._sparse:
-            row = row.toarray().ravel()
-        else:
-            row = np.asarray(row).reshape(-1)
-        return torch.from_numpy(row.astype(np.float32, copy=False))
-
-
-# ---------- 3. NPZ cell dataset (pretrain data) ----------
-
-class NpzCellDataset(Dataset):
-    """Lazily load per-cell npz files from a directory.
-
-    Supports two naming conventions:
-        * sequential: cell0.npz, cell1.npz, ... (set ``cell_num`` to total count)
-        * listing:    scan directory for all *.npz files (slower init, no cell_num needed)
-    """
-
-    def __init__(
-        self,
-        data_dir: str,
-        cell_num: Optional[int] = None,
-        max_cells: Optional[int] = None,
-        seed: int = 0,
-    ):
-        self.data_dir = data_dir
-        if cell_num is not None and cell_num > 0:
-            self.sequential = True
-            total = cell_num
-        else:
-            self.sequential = False
-            self._files = sorted(
-                f for f in os.listdir(data_dir) if f.endswith(".npz")
-            )
-            total = len(self._files)
-
-        if max_cells is not None and max_cells < total:
-            rng = np.random.default_rng(seed)
-            self.indices = np.sort(rng.choice(total, size=max_cells, replace=False))
-        else:
-            self.indices = np.arange(total)
-
-    def __len__(self) -> int:
-        return len(self.indices)
-
-    def __getitem__(self, idx: int) -> torch.Tensor:
-        real = int(self.indices[idx])
-        if self.sequential:
-            path = os.path.join(self.data_dir, f"cell{real}.npz")
-        else:
-            path = os.path.join(self.data_dir, self._files[real])
-        sp = sparse.load_npz(path)
-        arr = sp.toarray().squeeze() if sparse.issparse(sp) else np.asarray(sp).squeeze()
-        return torch.from_numpy(arr.astype(np.float32, copy=False))
-
-
-# ---------- 4. Misc ----------
-
-def ensure_dir(path: str) -> Path:
-    p = Path(path)
-    p.mkdir(parents=True, exist_ok=True)
-    return p

README.md

@@ -12,8 +12,7 @@ library_name: pytorch
 
 BioMod is a single-cell transcriptomics foundation model built around gene modules
 and an attention-based encoder. This repository provides pretrained model assets,
-pretraining and fine-tuning code, downstream task tutorials, and interpretability
-tools.
+pretraining and fine-tuning code, and downstream task tutorials.
 
 BioMod is released by **CTC-Atlas**, the joint laboratory of **Raymind Biotechnology
 Co., Ltd** and the **Shenzhen Research Institute of Big Data (SRIBD)**.
@@ -88,14 +87,13 @@ to end with the bundled checkpoint.
 | Drug-pair synergy classification | [Training script](BioMod/drug_combination/train_drug_combination_biomod.py) | [Notebook](BioMod/toturials/drug_combination.ipynb) |
 | Gene perturbation response prediction | [Training script](BioMod/gene_perturbation/train_biomod.py) | [Notebook](BioMod/toturials/gene_perturbation.ipynb) |
 
-### Pretraining and analysis scripts
+### Pretraining
 
-The repository also includes the following supporting workflows.
+The repository also supports pretraining and continued pretraining.
 
 | Workflow | Implementation | Documentation |
 | --- | --- | --- |
 | Pretraining and continued pretraining | [Pretraining script](BioMod/biomod.py) | [Source guide](BioMod/README.md), [launcher reference](BioMod/command.txt) |
-| Module interpretability analysis | [Module geometry](BioMod/interpretability/biomod_geometry.py), [module importance](BioMod/interpretability/biomod_importance/biomod_importance.py), [text alignment](BioMod/interpretability/biomod_text_alignment/biomod_text_alignment.py) | [Overview](BioMod/interpretability/README.md), [attention](BioMod/interpretability/attention/README.md), [module importance guide](BioMod/interpretability/biomod_importance/README.md) |
 
 Bundled input defaults resolve from each script's location. Command examples
 state their working directory; explicit relative command-line paths are resolved
@@ -139,7 +137,7 @@ Output directories are created when the corresponding operation is enabled.
 | --- | --- | --- |
 | Drug response / drug combination | `--checkpoint_path` | `--hyper_params_path` |
 | Zero-shot / gene perturbation | `--scfm_ckpt_path` | `--scfm_hyper_params_path` |
-| Cell annotation / backbone interpretation | `--biomod_checkpoint` | `--config_pkl` |
+| Cell annotation | `--biomod_checkpoint` | `--config_pkl` |
 | Continued pretraining | `--model_path` | `--model_config_path` or paired-file discovery |
 
 ### Data and evaluation
@@ -173,7 +171,7 @@ readability. Folder descriptions explain where to find each component.
 
 ```text
 BioMod/                                                   # Project root: source code, model dependencies and weights
-├── BioMod/                                               # Task code: training, evaluation, interpretability and tutorials
+├── BioMod/                                               # Task code: training, evaluation and tutorials
 │   ├── cell_annotation/                                  # Cell-type annotation and model fine-tuning
 │   │   ├── biomod_data/                                  # Hallmark and Reactome modules with loading utilities
 │   │   │   ├── load_biomod.py
@@ -197,41 +195,6 @@ BioMod/                                                   # Project root: source
 │   │   ├── scfm_utils_local.py
 │   │   ├── train_biomod.py
 │   │   └── utils.py
-│   ├── interpretability/                                 # Module representations and biological interpretability
-│   │   ├── attention/                                    # Module crosstalk and within-module gene attention
-│   │   │   ├── __init__.py
-│   │   │   ├── biomod_crosstalk.py
-│   │   │   ├── biomod_gene_attention.py
-│   │   │   ├── README.md
-│   │   │   └── README_gene_attention.md
-│   │   ├── biomod_importance/                            # Module importance, ablation studies and marker-gene validation
-│   │   │   ├── __init__.py
-│   │   │   ├── cross_class_similarity.py
-│   │   │   ├── family_ablation.py
-│   │   │   ├── family_importance_analysis.py
-│   │   │   ├── marker_gene_verification.py
-│   │   │   ├── masking_ablation.py
-│   │   │   ├── biomod_importance.py
-│   │   │   ├── biomod_importance_publication.py
-│   │   │   ├── pub_utils.py
-│   │   │   ├── random_baseline.py
-│   │   │   ├── README.md
-│   │   │   ├── RESULTS.md
-│   │   │   ├── tcell_reactome_analysis.py
-│   │   │   └── TODO.md
-│   │   ├── biomod_text_alignment/                        # Semantic alignment of module representations and pathway text
-│   │   │   ├── __init__.py
-│   │   │   ├── DESIGN.md
-│   │   │   ├── fetch_reactome_descriptions.py
-│   │   │   ├── MANUSCRIPT_SECTION.md
-│   │   │   ├── biomod_text_alignment.py
-│   │   │   ├── biomod_text_alignment_extended.py
-│   │   │   ├── biomod_text_alignment_publication.py
-│   │   │   └── reactome_descriptions.json
-│   │   ├── __init__.py
-│   │   ├── biomod_geometry.py
-│   │   ├── README.md
-│   │   └── utils.py
 │   ├── tests/                                            # CPU tests for shared training utilities
 │   │   └── test_training_utils.py
 │   ├── toturials/                                        # Jupyter tutorials for five downstream tasks