Upload BioMod with datasets and the 0810 checkpoint

Publish the BioMod source tree, model dependencies, complete downstream datasets, and the unchanged 0810 weight/configuration pair. Include the updated model implementation, shared training utilities, CPU utility tests, and English documentation. Exclude download caches, local environments, and OS metadata.

.gitattributes

@@ -33,3 +33,27 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/single_drug_response/data/GDSC/drug_graph_feat/5460769.hkl filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/drug_combination_synergy/data_summary/test_data/41467_2019_9799_MOESM3_ESM.xlsx filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/norman/perturb_processed.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/human_pancreas_norm_complexBatch_with_gene_name.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/single_drug_response/optional_reference/CCLE/genomic_mutation_34673_demap_features.csv filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/single_drug_response/data/GDSC/drug_graph_feat/6918289.hkl filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/Zheng68K.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/Kim2020_Lung_noMalignant.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/norman/go.csv filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/Myeloid/query_adata.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/Kim2020_Lung.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/BMMC_processed.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/Myeloid/reference_adata.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/MS/c_data.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/MS/filtered_ms_adata.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/hp14/demo_test.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/hp14/demo_train_with_MHC_class_II.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/drug_response_pytorch/drug_response_pytorch/DeepCDR/DeepCDR_ECCB_presentation.pdf filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/drug_response_pytorch/drug_response_pytorch/data/CCLE/genomic_mutation_34673_demap_features.csv filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/drug_response_pytorch/drug_response_pytorch/data/GDSC/drug_graph_feat/6918289.hkl filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/liver.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/drug_response_pytorch/drug_response_pytorch/data/GDSC/drug_graph_feat/5460769.hkl filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/Immune_ALL_human.h5ad filter=lfs diff=lfs merge=lfs -text
+data/downstream/final_data/drug_response_pytorch/drug_combination/data_summary/test_data/41467_2019_9799_MOESM3_ESM.xlsx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,42 @@
+# ===== 大文件：不进 GitHub，单独托管（HF / OSS）=====
+# weight/  4.1GB 训练权重 checkpoint
+# data/    17GB 下游公开 benchmark 数据集
+weight/
+data/
+!data/
+!data/downstream/
+!data/downstream/final_data/
+!data/downstream/final_data/norman/
+!data/downstream/final_data/norman/data_pyg/
+!data/downstream/final_data/norman/data_pyg/cell_graphs.pkl
+!data/downstream/final_data/norman/splits/
+!data/downstream/final_data/norman/splits/*.pkl
+
+# ===== 二进制 / 打包 / 缓存 =====
+*.zip
+!data/downstream/final_data/drug_response_pytorch/drug_response_pytorch/data/GDSC/drug_graph_feat.zip
+!data/downstream/final_data/single_drug_response/optional_reference/GDSC/drug_graph_feat.zip
+*.tar
+*.tar.gz
+*.pth
+*.pkl
+!data/downstream/final_data/norman/data_pyg/cell_graphs.pkl
+!data/downstream/final_data/norman/splits/*.pkl
+*.ckpt
+__pycache__/
+*.pyc
+.ipynb_checkpoints/
+
+# ===== 本地环境与密钥 =====
+.env
+*.local
+
+# Dated release checkpoints are intentionally versioned on Hugging Face.
+!weight/
+weight/*
+!weight/0629model/
+!weight/0629model/*.pkl
+!weight/0629model/*.pth
+!weight/0810model/
+!weight/0810model/*.pkl
+!weight/0810model/*.pth

BioMod/.env.example

@@ -0,0 +1,13 @@
+# Copy this file to `.env` or export the variables in your shell.
+# Never commit real API keys.
+
+# Alibaba DashScope OpenAI-compatible API, used by Qwen embedding/chat utilities.
+DASHSCOPE_API_KEY=
+DASHSCOPE_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1
+
+# Optional OpenAI-compatible DeepSeek endpoint used by archived text-quality scripts.
+DEEPSEEK_API_KEY=
+DEEPSEEK_BASE_URL=
+
+# Optional experiment tracking.
+WANDB_API_KEY=

BioMod/.gitattributes

@@ -0,0 +1,24 @@
+* text=auto eol=lf
+
+*.py text diff=python
+*.md text
+*.txt text
+*.yml text
+*.yaml text
+*.json text
+*.csv text
+*.sh text eol=lf
+*.ipynb text
+
+# Keep future large scientific artifacts out of normal Git history.
+# These patterns are compatible with Git LFS if release artifacts are needed.
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.h5ad filter=lfs diff=lfs merge=lfs -text
+*.loom filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text

BioMod/.gitignore

@@ -0,0 +1,56 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Local environments and secrets
+.env
+.env.*
+!.env.example
+.venv/
+venv/
+env/
+*.local
+
+# Logs and temporary outputs
+*.log
+logs/
+tmp/
+temp/
+scratch/
+nohup.out
+
+# Large datasets, checkpoints, and generated arrays
+*.h5
+*.h5ad
+*.loom
+*.zarr/
+*.npy
+*.npz
+*.pt
+*.pth
+*.pkl
+*.ckpt
+*.safetensors
+*.hkl
+*.hickle
+
+# Generated task outputs and caches
+gene_inverse_norman/outputs/
+toturials/execution_checks/
+stage2_alignment/archive_text_quality_eval/qwen_cache/
+stage2_alignment/best_scheme_qwen/cached_text_embeddings.pt
+
+# OS/editor files
+.DS_Store
+Thumbs.db
+.vscode/
+.idea/

BioMod/README.md

@@ -0,0 +1,428 @@
+# BioMod / PerformerModule
+
+Source code and tutorial notebooks for BioMod single-cell pretraining and
+downstream experiments using PerformerModule representations.
+
+This source directory contains the pretraining entry point and shared training
+utilities, alongside downstream task runners, tutorial notebooks, environment
+files, and plotting helpers.
+
+## Highlights
+
+| Area | What is included |
+| --- | --- |
+| Representation | Module-based single-cell encoder utilities and module metadata |
+| Annotation | Supervised cell-type annotation workflows |
+| Integration | Zero-shot batch integration and UMAP/scIB-style reporting |
+| Perturbation | Norman gene-perturbation tutorial and smoke workflow |
+| Drug response | Single-drug response prediction with GDSC/CCLE data |
+| Drug combinations | Drug-combination synergy prediction with independent-set evaluation |
+| Figures | Notebook-level PDF/PNG/SVG plotting cells for paper and GitHub display |
+
+## Repository Layout
+
+This is the `BioMod/` source directory inside the project root. See the
+[full project layout](../README.md#repository-layout) for sibling
+`model_dependencies/` and `weight/` directories; `data/` and `.cache/` are omitted
+from these repository trees.
+
+```text
+BioMod/                                                   # Task code: training, evaluation, interpretability and tutorials
+├── cell_annotation/                                      # Cell-type annotation and model fine-tuning
+│   ├── module_data/                                      # Hallmark and Reactome modules with loading utilities
+│   │   ├── load_module.py
+│   │   ├── module1_hallmarks.json
+│   │   └── module2_reactome_2022_filtered.json
+│   ├── cmd_finetune.md
+│   └── performer_module_finetune.py
+├── drug_combination/                                     # Drug-combination synergy prediction
+│   ├── README.md
+│   └── train_drug_combination_performer_module.py
+├── drug_response/                                        # Single-drug response prediction and training
+│   ├── cmd.txt
+│   └── train_drug_performer_module.py
+├── eval_new_data/                                        # Reconstruction error and generalization on new data
+│   ├── cmd.txtx
+│   └── evaluate_new_data.py
+├── gene_inverse_norman/                                  # Norman reverse perturbation prediction and evaluation
+│   ├── protocols/                                        # Preset gene sets and train/validation/test splits
+│   │   ├── cellfm/                                       # Gene list and data splits for the CellFM protocol
+│   │   │   ├── genes.txt
+│   │   │   └── split.json
+│   │   └── scgpt/                                        # Gene list and data splits for the scGPT protocol
+│   │       ├── genes.txt
+│   │       └── split.json
+│   ├── build_norman_gene_inverse_benchmark.py
+│   ├── common.py
+│   ├── evaluate_gene_inverse.py
+│   ├── finetune_gene_inverse.py
+│   ├── README.md
+│   └── utils.py
+├── gene_perturbation/                                    # Perturbation response prediction: processing, training and inference
+│   ├── data_utils_local.py
+│   ├── gears_reduce_local.py
+│   ├── inference_local.py
+│   ├── model_module.py
+│   ├── pertdata_local.py
+│   ├── README.md
+│   ├── scfm_utils_local.py
+│   ├── train_module.py
+│   └── utils.py
+├── interpretability/                                     # Module representations and biological interpretability
+│   ├── attention/                                        # Module crosstalk and within-module gene attention
+│   │   ├── __init__.py
+│   │   ├── module_crosstalk.py
+│   │   ├── module_gene_attention.py
+│   │   ├── README.md
+│   │   └── README_gene_attention.md
+│   ├── module_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
+│   │   ├── module_importance.py
+│   │   ├── module_importance_publication.py
+│   │   ├── pub_utils.py
+│   │   ├── random_baseline.py
+│   │   ├── README.md
+│   │   ├── RESULTS.md
+│   │   ├── tcell_reactome_analysis.py
+│   │   └── TODO.md
+│   ├── module_text_alignment/                            # Semantic alignment of module representations and pathway text
+│   │   ├── __init__.py
+│   │   ├── DESIGN.md
+│   │   ├── fetch_reactome_descriptions.py
+│   │   ├── MANUSCRIPT_SECTION.md
+│   │   ├── module_text_alignment.py
+│   │   ├── module_text_alignment_extended.py
+│   │   ├── module_text_alignment_publication.py
+│   │   └── reactome_descriptions.json
+│   ├── __init__.py
+│   ├── module_geometry.py
+│   ├── README.md
+│   └── utils.py
+├── module_data_all/                                      # Module definitions, merging and uncovered-gene tools
+│   ├── create_uncovered_modules.py
+│   ├── merge_modules.py
+│   └── merged_modules.json
+├── tests/                                                # CPU tests for shared training utilities
+│   └── test_training_utils.py
+├── toturials/                                            # Jupyter tutorials for five downstream tasks
+│   ├── cell_annotation.ipynb
+│   ├── drug_combination.ipynb
+│   ├── drug_response_prediction.ipynb
+│   ├── gene_perturbation.ipynb
+│   └── zero_shot_batch_integration.ipynb
+├── zero_shot_batch/                                      # Zero-shot batch integration, evaluation and visualization
+│   ├── sequencing_depth/                                 # Low-depth data generation and zero-shot evaluation
+│   │   ├── cmd.txt
+│   │   ├── make_low_depth_h5ad.py
+│   │   └── run_zero_shot_low_depth.py
+│   ├── performer_module_zero_shot.py
+│   ├── README.md
+│   └── visualize_results.py
+├── .env.example
+├── .gitattributes
+├── .gitignore
+├── command.txt
+├── environment.yml
+├── models_performer_module.py
+├── module_dataset.py
+├── performer_module.py
+├── README.md
+├── requirements.txt
+└── training_utils.py
+```
+
+The notebook directory is intentionally named `toturials` to preserve existing
+paths used by the project.
+
+## Quick Start
+
+Create a fresh environment:
+
+```bash
+cd /path/to/BioMod/BioMod
+conda env create -f environment.yml
+conda activate performer_module
+```
+
+Or install Python packages into an existing CUDA/PyTorch environment:
+
+```bash
+python -m pip install -r requirements.txt
+```
+
+For the original server runs, the tutorials assume the `sclong` environment:
+
+```bash
+conda activate sclong
+jupyter lab /path/to/BioMod/BioMod/toturials
+```
+
+### Training utility checks
+
+The pretraining entry point imports `training_utils.py` explicitly. This local
+module supplies random seeding, a sequential distributed validation sampler,
+cosine learning-rate scheduling with warmup and restarts, scalar metric
+reduction, and resumable checkpoint saving. It does not depend on a separate
+package named `utils` or on a downstream task's utility module.
+
+From this source directory, run the CPU utility tests and check entry-point
+imports after installing the dependencies:
+
+```bash
+python -B -m unittest discover -s tests -p test_training_utils.py -v
+python -B performer_module.py --help
+```
+
+The utility tests use synthetic tensors and temporary checkpoints; they do not
+load the supplied pretrained weights. Passing these checks does not validate a
+full training run. The pretraining entry point still requires CUDA and an NCCL
+process group, with training data and paths configured for the target machine.
+
+## Shared Assets
+
+This local checkout retains only `weight/0810model/`. The original 0629
+checkpoint/config pair is not included. For examples that reference
+`weight/0629model/`, use the matching 0810 `.pth` and `.pkl` files described below,
+and update any server-specific paths for the local environment.
+
+Tutorials expect a pretrained BioMod/PerformerModule checkpoint and shared
+gene assets. Set `BIOMODULEFM_ASSET_ROOT` to the directory containing
+`weight/0810model/` and `model_dependencies/` (the outer `BioMod/` project directory).
+
+```bash
+export BIOMODULEFM_ASSET_ROOT=/path/to/BioMod
+```
+
+Current local asset files:
+
+```text
+$BIOMODULEFM_ASSET_ROOT/weight/0810model/                # 0810 checkpoint: .pth weights and .pkl configuration
+├── mae_tost_70M_75mask_ec2.pkl
+└── mae_tost_70M_75mask_ec2_epoch0.pth
+
+$BIOMODULEFM_ASSET_ROOT/model_dependencies/              # Shared gene mappings, gene lists, embeddings and model components
+├── module_data_all/                                     # Shared module definitions and construction tools
+│   ├── create_uncovered_modules.py
+│   ├── merge_modules.py
+│   └── merged_modules.json
+├── performer_pytorch_cont/                              # Performer attention and model components
+│   ├── performer_pytorch_cont/                          # Nested package containing copies of the model component source files
+│   │   ├── __init__.py
+│   │   ├── ding_models.py
+│   │   ├── performer_pytorch.py
+│   │   ├── performer_pytorch_old.py
+│   │   └── reversible.py
+│   ├── __init__.py
+│   ├── ding_models.py
+│   ├── performer_pytorch.py
+│   ├── performer_pytorch_old.py
+│   └── reversible.py
+├── gene2ens.csv
+├── selected_gene2vec_27k.npy
+└── selected_genes_27k.txt
+```
+
+Project-local module definitions are kept in:
+
+```text
+BioMod/module_data_all/                                  # Module definitions, merging and uncovered-gene tools
+├── create_uncovered_modules.py
+├── merge_modules.py
+└── merged_modules.json
+```
+
+## Smaller 0810 model
+
+The repository also includes a smaller 2-layer encoder checkpoint:
+
+```text
+weight/0810model/                                        # 0810 checkpoint: .pth weights and .pkl configuration
+|-- mae_tost_70M_75mask_ec2_epoch0.pth
+`-- mae_tost_70M_75mask_ec2.pkl
+```
+
+It uses a 2-layer encoder with the same 27,874-gene axis and 660 module
+definitions as the primary 32-layer model. Use its checkpoint and PKL together
+by replacing the two `weight/0629model/` paths in the commands below with the
+matching `weight/0810model/` files. The code restores `enc_depth=2` and the
+other architecture settings from the PKL automatically.
+
+## Tutorials
+
+| Task | Notebook | Main entry point | Default output style |
+| --- | --- | --- | --- |
+| Cell-type annotation | `toturials/cell_annotation.ipynb` | task script under `cell_annotation/` | metrics, classification report, confusion matrix |
+| Zero-shot batch integration | `toturials/zero_shot_batch_integration.ipynb` | `zero_shot_batch/` runner | UMAP PDF/PNG, avgBatch, avgBio |
+| Gene perturbation | `toturials/gene_perturbation.ipynb` | `gene_perturbation/` runner | GEARS-style metrics and smoke outputs |
+| Drug response prediction | `toturials/drug_response_prediction.ipynb` | `drug_response/train_drug_performer_module.py` | scatter, residuals, ROC, per-drug and per-cancer tables |
+| Drug-combination synergy | `toturials/drug_combination.ipynb` | `drug_combination/train_drug_combination_performer_module.py` | label balance, loss, AUROC, PR-AUC, confusion matrix when predictions exist |
+
+Each notebook uses absolute paths through configurable roots and keeps full
+training disabled by default when data are not present. Flip the notebook
+`RUN_TRAINING` flag only after checking the path table at the top of the
+notebook.
+
+## Data Layout
+
+Large datasets are not committed to this code repository. Place them under
+task-specific data directories or pass absolute paths through environment
+variables.
+
+### Single-Drug Response
+
+Expected clean data root:
+
+```text
+single_drug_response/data/
+|-- CCLE/
+|   |-- Cell_lines_annotations_20181226.txt
+|   |-- GDSC_IC50.csv
+|   `-- genomic_expression_561celllines_697genes_demap_features.csv
+`-- GDSC/
+    |-- 1.Drug_listMon Jun 24 09_00_55 2019.csv
+    `-- drug_graph_feat/
+```
+
+### Drug-Combination Synergy
+
+Expected clean data root:
+
+```text
+drug_combination/data_summary/
+|-- new_labels_0_10.csv
+|-- new_cell_features_954.csv
+|-- smiles_0_10.csv
+`-- test_data/
+    |-- independent_input.csv
+    `-- independent_cell_features_954.csv
+```
+
+The drug-combination tutorial also supports old DeepDDS-style prediction CSV
+files for ROC, PR, and confusion-matrix visualization.
+
+## Example Commands
+
+Drug response prediction with the primary 0629 model:
+
+```bash
+python drug_response/train_drug_performer_module.py \
+  --gpu_id 0 \
+  --data_path /path/to/single_drug_response/data \
+  --checkpoint_path $BIOMODULEFM_ASSET_ROOT/weight/0629model/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+  --hyper_params_path $BIOMODULEFM_ASSET_ROOT/weight/0629model/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+  --module_data_dir ./module_data_all \
+  --gene_list_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/selected_genes_27k.txt \
+  --gene_mapping_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/gene2ens.csv \
+  --gene2vec_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/selected_gene2vec_27k.npy \
+  --symbol2ens_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/gene2ens.csv \
+  --save_dir ./results/drug_response_0629 \
+  --epochs 300 \
+  --batch_size 32 \
+  --pool_method max
+```
+
+Drug-combination synergy prediction:
+
+```bash
+python drug_combination/train_drug_combination_performer_module.py \
+  --checkpoint_path $BIOMODULEFM_ASSET_ROOT/weight/0629model/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+  --hyper_params_path $BIOMODULEFM_ASSET_ROOT/weight/0629model/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+  --module_data_dir ./module_data_all \
+  --gene_list_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/selected_genes_27k.txt \
+  --gene_mapping_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/gene2ens.csv \
+  --gene2vec_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/selected_gene2vec_27k.npy \
+  --symbol2ens_path $BIOMODULEFM_ASSET_ROOT/model_dependencies/gene2ens.csv \
+  --data_dir ./drug_combination/data_summary \
+  --output_dir ./results/drug_combination_0629 \
+  --epochs 500 \
+  --batch_size 256 \
+  --lr 1e-4 \
+  --n_folds 1 \
+  --pool_method max
+```
+
+## Display Outputs
+
+The notebooks regenerate figure files in task output directories. Typical files
+include:
+
+| Task | Generated display files |
+| --- | --- |
+| Zero-shot integration | A4 UMAP PDF/PNG panels, combined cell-type and batch views |
+| Drug response | `drug_response_overall_scatter.pdf`, `drug_response_median_threshold_roc.svg`, per-drug and per-cancer CSV summaries |
+| Drug combination | `drug_combination_metric_curves.pdf`, `drug_combination_roc_curves.pdf`, `drug_combination_label_balance.pdf`, optional PR and confusion-matrix panels |
+| Cell annotation | confusion matrices, classification reports, UMAP/prediction diagnostics |
+| Gene perturbation | smoke-run logs, perturbation metrics, downstream result tables |
+
+For GitHub display, use PNG files. For manuscripts, use the PDF/SVG exports.
+The plotting cells use embedded, editable PDF fonts where supported.
+
+## Reference Results
+
+These values are retained as run-level reference points and should be refreshed
+from final locked runs before manuscript submission.
+
+| Task | Dataset / setting | Reported signal |
+| --- | --- | --- |
+| Drug response | GDSC/CCLE, primary 0629 model | Pearson correlation around `0.93` in the existing full run |
+| Drug combination | Independent-set domain generalization, primary 0629 model | AUROC reference around `0.6824` for the previous run |
+| Zero-shot integration | HP, BMMC, Kim Lung, Kim Lung no-malignant | Notebook computes avgBatch and avgBio across repeated runs |
+| Cell annotation | Seven benchmark datasets | Notebook reports accuracy, macro-F1, and confusion matrices |
+
+## Reproducibility Conventions
+
+- Record the exact command line, checkpoint, hyper-parameter pkl, data version,
+  random seed, and output directory for every final run.
+- Keep raw data, cached embeddings, model checkpoints, metric tables, and paper
+  figures in separate directories.
+- Prefer shared storage for large `.h5ad`, `.npy`, and checkpoint outputs.
+- Do not commit private data, raw model weights, API keys, or personal absolute
+  paths.
+
+## Troubleshooting
+
+### Missing data
+
+Open the relevant notebook and run the path-check cell. It prints all required
+files and shows which entries are missing.
+
+### Missing model assets
+
+Check `BIOMODULEFM_ASSET_ROOT` and verify that the 0407 checkpoint, pkl,
+gene2vec, gene list, gene mapping, and `performer_pytorch_cont/` directory
+exist.
+
+### No space left on device
+
+Large `.h5ad`, `.npy`, checkpoint, and figure outputs can fill local disks.
+Move result directories to shared storage and remove partial failed outputs only
+after confirming they are not needed.
+
+### Matplotlib font warnings
+
+The notebooks request Arial first and fall back to DejaVu Sans when Arial is not
+installed. PDF font embedding is enabled through `pdf.fonttype = 42`.
+
+## GitHub Release Checklist
+
+- Keep `README.md`, `requirements.txt`, `environment.yml`, `.gitignore`, and
+  `.gitattributes` at the repository root.
+- Confirm tutorial notebooks open without personal paths.
+- Keep large data and checkpoints outside Git, then document download or
+  placement instructions.
+- Add final figure examples under a small `docs/` or `assets/` directory if the
+  repository needs rendered images on the GitHub landing page.
+- Add license, citation, and data-availability text before public release.
+
+## Changelog
+
+| Date | Change |
+| --- | --- |
+| 2026-06-15 | Added paper-style plotting cells for single-drug and drug-combination tutorials. |
+| 2026-06-15 | Removed `stage2_alignment/` from the active repository scope. |
+| 2026-06-15 | Added root environment files and GitHub-ready repository documentation. |

BioMod/cell_annotation/cmd_finetune.md

@@ -0,0 +1,15 @@
+torchrun --nproc_per_node=4 performer_module_finetune.py \
+--config_pkl ../ckpts_performer_module_multiply_all_module/performer_module_multiply_all_module.pkl \
+--performer_checkpoint ../ckpts_performer_module_multiply_all_module/performer_module_multiply_all_module_epoch6.pth \
+--data_path ../../cell_annotation/Zheng68K.h5ad \
+--label_column celltype \
+--head_type attn_pool \
+--finetune_mode last_n \
+--unfreeze_last_n 1 \
+--cls_decoder_layers 3 \
+--lr_head 1e-4 \
+--lr_backbone 1e-5 \
+--batch_size 256 \
+--epochs 300 \
+--module_data_dir ../module_data_all \
+--output_dir ../../../../share/extracted_feature/ckpts_performer_module_fintune/performer_module_multiply_all_module

BioMod/cell_annotation/module_data/load_module.py

@@ -0,0 +1,33 @@
+def load_json_files(file1_path, file2_path):
+    """
+    
+    Args:
+        file1_path: hallmarks file path
+        file2_path: reactome file path
+        
+    Returns:
+        tuple: (hallmark_dict, filtered_reactome_dict)
+    """
+    # Read HALLMARKS
+    with open(file1_path, 'r', encoding='utf-8') as f:
+        hallmark_dict = json.load(f)
+    
+    # Read reactome data
+    with open(file2_path, 'r', encoding='utf-8') as f:
+        reactome_full = json.load(f)
+    
+    # Filter reactome data
+    filtered_reactome = {
+        k: v for k, v in reactome_full.items() 
+        if 50 <= len(v) <= 500
+    }
+    
+    return hallmark_dict, filtered_reactome
+
+
+# hallmark_data, reactome_data = load_json_files(
+#     'module1_hallmarks.json', 
+#     'module2_reactome.json'
+# )
+# len(hallmark_data.keys())
+# len(reactome_data.keys())

BioMod/cell_annotation/performer_module_finetune.py

@@ -0,0 +1,1591 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+End-to-end cell annotation fine-tuning for Performer Module.
+
+This script is intentionally independent from the legacy two-stage pipeline:
+1. No feature extraction to .npy files.
+2. No imports from other scripts in performer_module/cell_annotation.
+3. Reuses the pretrained Performer Module backbone directly.
+
+Supported cell embedding heads:
+- module_mean: mean over restored module tokens
+- max: max pooling over restored module tokens
+- gene_mean: module -> gene -> mean over genes
+- attn_pool: learnable attention pooling over restored module tokens
+- cls: use pretrained CLS token if the checkpoint contains one
+
+Supported backbone finetune modes:
+- frozen: freeze the full backbone
+- last_n: unfreeze only the last N encoder blocks
+- full: full-parameter finetuning
+"""
+
+import argparse
+import contextlib
+import json
+import logging
+import os
+import pickle
+import random
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+import anndata as ad
+import torch.distributed as dist
+from scipy import sparse
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score
+from sklearn.model_selection import StratifiedShuffleSplit
+from sklearn.preprocessing import LabelEncoder
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch import nn
+from torch.utils.data import DataLoader, Dataset, Sampler
+from torch.utils.data.distributed import DistributedSampler
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+PERFORMER_MODULE_ROOT = SCRIPT_DIR.parent
+SCLONG_ROOT = PERFORMER_MODULE_ROOT.parent
+
+if str(PERFORMER_MODULE_ROOT) not in sys.path:
+    sys.path.insert(0, str(PERFORMER_MODULE_ROOT))
+if str(SCLONG_ROOT) not in sys.path:
+    sys.path.insert(0, str(SCLONG_ROOT))
+
+from models_performer_module import (
+    migrate_legacy_module_config,
+    migrate_legacy_module_state_dict,
+)
+
+if TYPE_CHECKING:
+    from models_performer_module import PerformerEncoderOnly_Module
+
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def is_dist_initialized() -> bool:
+    return dist.is_available() and dist.is_initialized()
+
+
+def get_rank() -> int:
+    return dist.get_rank() if is_dist_initialized() else 0
+
+
+def get_world_size() -> int:
+    return dist.get_world_size() if is_dist_initialized() else 1
+
+
+def is_main_process() -> bool:
+    return get_rank() == 0
+
+
+def unwrap_model(model: nn.Module) -> nn.Module:
+    return model.module if hasattr(model, "module") else model
+
+
+def init_distributed_mode(device_arg: str) -> Tuple[torch.device, bool, int, int, int]:
+    use_ddp = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    if not use_ddp:
+        if device_arg == "cuda" and not torch.cuda.is_available():
+            logger.warning("CUDA requested but not available. Falling back to CPU.")
+            return torch.device("cpu"), False, 0, 1, 0
+        return torch.device(device_arg), False, 0, 1, 0
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+
+    if device_arg != "cuda":
+        backend = "gloo"
+        device = torch.device(device_arg)
+    else:
+        if not torch.cuda.is_available():
+            raise RuntimeError("DDP with device='cuda' requires CUDA to be available")
+        backend = "nccl"
+        torch.cuda.set_device(local_rank)
+        device = torch.device("cuda", local_rank)
+
+    dist.init_process_group(backend=backend)
+    return device, True, rank, world_size, local_rank
+
+
+def all_gather_numpy(array: np.ndarray) -> np.ndarray:
+    if not is_dist_initialized():
+        return array
+    gathered = [None for _ in range(get_world_size())]
+    dist.all_gather_object(gathered, array)
+    non_empty = [arr for arr in gathered if arr is not None and len(arr) > 0]
+    if not non_empty:
+        return np.empty((0,), dtype=array.dtype)
+    return np.concatenate(non_empty, axis=0)
+
+
+def reduce_sum(value: float, device: torch.device) -> float:
+    if not is_dist_initialized():
+        return float(value)
+    tensor = torch.tensor(float(value), device=device, dtype=torch.float64)
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    return float(tensor.item())
+
+
+def resolve_path(path_str: str) -> Path:
+    path = Path(path_str)
+    if path.is_absolute():
+        return path
+    return (SCRIPT_DIR / path).resolve()
+
+
+def seed_all(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+class DataPreprocessor:
+    """Map input genes to the pretrained model gene space."""
+
+    def __init__(
+        self,
+        gene2ens_file: str,
+        gene2ens_column: str = "Ensembl gene ID",
+        model_genes_file: Optional[str] = None,
+    ):
+        self.gene2ens_file = gene2ens_file
+        self.gene2ens_column = gene2ens_column
+        self.model_genes_file = model_genes_file
+        self.gene2ens_dict: Optional[Dict[str, str]] = None
+        self.model_genes: Optional[List[str]] = None
+        self._index_cache: Dict[str, torch.Tensor] = {}
+
+    def load_gene_mapping(self) -> None:
+        gene2ens = pd.read_csv(self.gene2ens_file, index_col=0)
+        if self.gene2ens_column not in gene2ens.columns:
+            raise ValueError(f"Column '{self.gene2ens_column}' not found in {self.gene2ens_file}")
+        self.gene2ens_dict = gene2ens[self.gene2ens_column].to_dict()
+        logger.info("Loaded gene2ens mapping: %d entries", len(self.gene2ens_dict))
+
+    def load_model_genes(self) -> None:
+        if self.model_genes_file is None:
+            raise ValueError("model_genes_file must be provided")
+        with open(self.model_genes_file, "r", encoding="utf-8") as handle:
+            self.model_genes = [line.rstrip("\n") for line in handle if line.strip()]
+        logger.info("Loaded %d model genes from %s", len(self.model_genes), self.model_genes_file)
+
+    def _extract_input_genes(
+        self,
+        adata: ad.AnnData,
+        gene_name_column: Optional[str] = None,
+    ) -> Tuple[List[str], str]:
+        if gene_name_column is not None and gene_name_column.lower() == "index":
+            return adata.var.index.astype(str).tolist(), "var.index"
+
+        if gene_name_column is not None and gene_name_column.lower() != "auto":
+            if gene_name_column not in adata.var.columns:
+                raise ValueError(f"Gene column '{gene_name_column}' not found in adata.var")
+            return adata.var[gene_name_column].astype(str).tolist(), f"var['{gene_name_column}']"
+
+        if "gene_name" in adata.var.columns:
+            return adata.var["gene_name"].astype(str).tolist(), "var['gene_name']"
+        return adata.var.index.astype(str).tolist(), "var.index"
+
+    def _build_reindex_state(
+        self,
+        adata: ad.AnnData,
+        dataset_name: str,
+        gene_name_column: Optional[str] = None,
+    ) -> dict:
+        input_genes, gene_source = self._extract_input_genes(adata, gene_name_column=gene_name_column)
+        if self.gene2ens_dict is not None:
+            input_genes = [self.gene2ens_dict.get(symbol, symbol) for symbol in input_genes]
+            logger.info("[%s] Mapped input genes to Ensembl IDs", dataset_name)
+
+        input_mapping = {gene: i for i, gene in enumerate(input_genes)}
+        input_gene_num = len(input_genes)
+
+        if self.model_genes is None:
+            raise RuntimeError("Model genes must be loaded before preprocess_adata")
+
+        model_index_positions = [input_mapping.get(gene, input_gene_num) for gene in self.model_genes]
+        overlap = sum(pos < input_gene_num for pos in model_index_positions)
+        logger.info(
+            "[%s] Gene source: %s | overlap with model space: %d / %d (model gene space: %d)",
+            dataset_name,
+            gene_source,
+            overlap,
+            input_gene_num,
+            len(self.model_genes),
+        )
+        return {
+            "cache_key": f"{dataset_name}:{gene_source}:{input_gene_num}",
+            "dataset_name": dataset_name,
+            "gene_source": gene_source,
+            "model_index_positions": model_index_positions,
+            "input_gene_num": input_gene_num,
+            "overlap": overlap,
+        }
+
+    def preprocess_adata(
+        self,
+        adata_path: str,
+        min_genes: int = 10,
+        min_cells: int = 10,
+        max_cells: Optional[int] = None,
+        dataset_name: str = "dataset",
+        gene_name_column: Optional[str] = None,
+    ) -> Tuple[ad.AnnData, dict]:
+        logger.info("Loading data from %s", adata_path)
+        adata = ad.read_h5ad(adata_path)
+        matrix = adata.X
+
+        if sparse.issparse(matrix):
+            cell_nonzero = np.asarray((matrix > 0).sum(axis=1)).ravel()
+            gene_nonzero = np.asarray((matrix > 0).sum(axis=0)).ravel()
+        else:
+            cell_nonzero = np.count_nonzero(np.asarray(matrix), axis=1)
+            gene_nonzero = np.count_nonzero(np.asarray(matrix), axis=0)
+
+        cell_mask = cell_nonzero >= min_genes
+        gene_mask = gene_nonzero >= min_cells
+        adata = adata[cell_mask, gene_mask].copy()
+
+        if max_cells is not None and adata.n_obs > max_cells:
+            logger.info("Subsampling %d cells from %d", max_cells, adata.n_obs)
+            rng = np.random.default_rng(0)
+            chosen = np.sort(rng.choice(adata.n_obs, size=max_cells, replace=False))
+            adata = adata[chosen].copy()
+
+        reindex_state = self._build_reindex_state(
+            adata,
+            dataset_name=dataset_name,
+            gene_name_column=gene_name_column,
+        )
+        return adata, reindex_state
+
+    def _get_index_tensor(self, reindex_state: dict, device: torch.device) -> torch.Tensor:
+        key = f"{str(device)}::{reindex_state['cache_key']}"
+        if key not in self._index_cache:
+            self._index_cache[key] = torch.as_tensor(
+                reindex_state["model_index_positions"],
+                dtype=torch.long,
+                device=device,
+            )
+        return self._index_cache[key]
+
+    def reindex_tensor_to_model_space(
+        self,
+        tensor: torch.Tensor,
+        reindex_state: dict,
+        filler: float = 0.0,
+    ) -> torch.Tensor:
+        if tensor.dim() != 2:
+            raise ValueError(f"Expected 2D tensor [B, G], got {tuple(tensor.shape)}")
+        index_tensor = self._get_index_tensor(reindex_state, tensor.device)
+        filler_col = torch.full(
+            (tensor.shape[0], 1),
+            float(filler),
+            dtype=tensor.dtype,
+            device=tensor.device,
+        )
+        expanded_tensor = torch.cat([tensor, filler_col], dim=1)
+        return torch.index_select(expanded_tensor, 1, index_tensor)
+
+
+class ExpressionDataset(Dataset):
+    def __init__(self, matrix, labels: np.ndarray, indices: Sequence[int]):
+        self.matrix = matrix.tocsr() if sparse.issparse(matrix) and not sparse.isspmatrix_csr(matrix) else matrix
+        self.labels = labels.astype(np.int64, copy=False)
+        self.indices = np.asarray(indices, dtype=np.int64)
+
+    def __len__(self) -> int:
+        return len(self.indices)
+
+    def __getitem__(self, idx: int):
+        real_idx = int(self.indices[idx])
+        row = self.matrix[real_idx]
+        if sparse.issparse(row):
+            row = row.toarray().ravel()
+        else:
+            row = np.asarray(row).reshape(-1)
+        x = torch.from_numpy(row.astype(np.float32, copy=False))
+        y = torch.tensor(self.labels[real_idx], dtype=torch.long)
+        return x, y
+
+
+class DistributedEvalSampler(Sampler[int]):
+    """Shard evaluation data across ranks without padding or duplication."""
+
+    def __init__(self, dataset_len: int, rank: int, world_size: int):
+        self.indices = list(range(rank, dataset_len, world_size))
+
+    def __iter__(self):
+        return iter(self.indices)
+
+    def __len__(self) -> int:
+        return len(self.indices)
+
+
+class GatedAttentionPool(nn.Module):
+    def __init__(self, embed_dim: int, hidden_dim: int):
+        super().__init__()
+        self.score = nn.Sequential(
+            nn.LayerNorm(embed_dim),
+            nn.Linear(embed_dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, 1),
+        )
+
+    def forward(self, tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        scores = self.score(tokens).squeeze(-1)
+        weights = torch.softmax(scores, dim=1)
+        pooled = torch.sum(tokens * weights.unsqueeze(-1), dim=1)
+        return pooled, weights
+
+
+class ScGPTClsDecoder(nn.Module):
+    """Classifier head aligned with scGPT's ClsDecoder."""
+
+    def __init__(self, d_model: int, n_cls: int, nlayers: int = 3):
+        super().__init__()
+        if nlayers < 1:
+            raise ValueError(f"nlayers must be >= 1, got {nlayers}")
+        self.layers = nn.ModuleList()
+        for _ in range(nlayers - 1):
+            self.layers.append(nn.Linear(d_model, d_model))
+            self.layers.append(nn.ReLU())
+            self.layers.append(nn.LayerNorm(d_model))
+        self.out_layer = nn.Linear(d_model, n_cls)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(x)
+        return self.out_layer(x)
+
+
+class ModuleSelector(nn.Module):
+    """Per-module scorer: soft weighting during training, hard threshold during inference."""
+
+    def __init__(self, embed_dim: int):
+        super().__init__()
+        self.scorer = nn.Sequential(
+            nn.LayerNorm(embed_dim),
+            nn.Linear(embed_dim, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(
+        self, x: torch.Tensor, tau: float = 0.5, hard: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        scores = self.scorer(x)  # (B, L, 1)
+        if hard:
+            x = x * (scores > tau).float()
+        else:
+            x = x * scores
+        return x, scores
+
+
+class PerformerModuleCellClassifier(nn.Module):
+    def __init__(
+        self,
+        backbone: "PerformerEncoderOnly_Module",
+        module_sets: List[List[int]],
+        num_classes: int,
+        head_type: str,
+        pool_hidden_dim: int,
+        dropout: float,
+        cls_decoder_layers: int,
+        use_module_selector: bool = False,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.module_sets = module_sets
+        self.num_classes = num_classes
+        self.head_type = head_type
+        self.base_module_count = len(module_sets)
+        self.embed_dim = backbone.base_dim
+        self.num_genes = backbone.num_genes
+        self.use_cls = bool(backbone.use_cls)
+
+        if self.head_type == "cls" and not self.use_cls:
+            raise ValueError("head_type='cls' requires a checkpoint with cls_token")
+
+        if self.head_type == "attn_pool":
+            hidden_dim = pool_hidden_dim if pool_hidden_dim > 0 else self.embed_dim
+            self.pooler = GatedAttentionPool(self.embed_dim, hidden_dim)
+        else:
+            self.pooler = None
+        self.output_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+        self.use_module_selector = use_module_selector
+        self.selector_tau = 0.5
+        if use_module_selector:
+            self.module_selector = ModuleSelector(self.embed_dim)
+        else:
+            self.module_selector = None
+        self._selector_scores = None
+
+        if self.head_type == "gene_mean":
+            module_idx, gene_idx, gene_counts = self._build_module_gene_index()
+            self.register_buffer("agg_module_idx", module_idx)
+            self.register_buffer("agg_gene_idx", gene_idx)
+            self.register_buffer("agg_gene_counts", gene_counts)
+        else:
+            self.agg_module_idx = None
+            self.agg_gene_idx = None
+            self.agg_gene_counts = None
+
+        self.classifier = ScGPTClsDecoder(
+            d_model=self.embed_dim,
+            n_cls=num_classes,
+            nlayers=cls_decoder_layers,
+        )
+
+    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
+        """Accept current module names and older serialized backbone names."""
+        return super().load_state_dict(
+            migrate_legacy_module_state_dict(state_dict),
+            strict=strict,
+            assign=assign,
+        )
+
+    def _build_module_gene_index(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        module_indices = []
+        gene_indices = []
+        gene_counts = torch.zeros(self.num_genes, dtype=torch.long)
+        for module_idx, module_genes in enumerate(self.module_sets):
+            for gene_idx in module_genes:
+                if 0 <= gene_idx < self.num_genes:
+                    module_indices.append(module_idx)
+                    gene_indices.append(gene_idx)
+                    gene_counts[gene_idx] += 1
+        return (
+            torch.tensor(module_indices, dtype=torch.long),
+            torch.tensor(gene_indices, dtype=torch.long),
+            gene_counts,
+        )
+
+    def _restore_module_order(self, latent: torch.Tensor, ids_restore: torch.Tensor) -> torch.Tensor:
+        if self.use_cls:
+            module_tokens = latent[:, 1:, :]
+        else:
+            module_tokens = latent
+        gather_index = ids_restore.unsqueeze(-1).expand(-1, -1, module_tokens.shape[-1])
+        return torch.gather(module_tokens, dim=1, index=gather_index)
+
+    def _module_to_gene(self, module_tokens: torch.Tensor) -> torch.Tensor:
+        batch_size, _, dim = module_tokens.shape
+        gene_feats = torch.zeros(
+            batch_size,
+            self.num_genes,
+            dim,
+            device=module_tokens.device,
+            dtype=module_tokens.dtype,
+        )
+        global_mean = module_tokens.mean(dim=1)
+
+        base_counts = self.agg_gene_counts.to(device=module_tokens.device, dtype=module_tokens.dtype)
+        gene_counts = base_counts.unsqueeze(0).expand(batch_size, -1).clone()
+
+        static_module_tokens = module_tokens[:, :self.base_module_count, :]
+        if self.agg_module_idx.numel() > 0:
+            pair_module_feats = static_module_tokens[:, self.agg_module_idx, :]
+            gene_idx_expand = self.agg_gene_idx.view(1, -1, 1).expand(batch_size, -1, dim)
+            gene_feats.scatter_add_(1, gene_idx_expand, pair_module_feats)
+
+        zero_mask = gene_counts == 0
+        gene_feats = gene_feats / gene_counts.clamp(min=1.0).unsqueeze(-1)
+        gene_feats = torch.where(zero_mask.unsqueeze(-1), global_mean.unsqueeze(1), gene_feats)
+        return gene_feats
+
+    def get_selector_sparsity_loss(self) -> torch.Tensor:
+        if self._selector_scores is not None:
+            return self._selector_scores.mean()
+        return torch.tensor(0.0)
+
+    def forward(self, x_model: torch.Tensor) -> torch.Tensor:
+        latent, _, ids_restore, _ = self.backbone.forward_encoder(x_model, mask_ratio=0.0)
+
+        if self.head_type == "cls":
+            cell_emb = latent[:, 0, :]
+        else:
+            module_tokens = self._restore_module_order(latent, ids_restore)
+            if self.module_selector is not None:
+                module_tokens, self._selector_scores = self.module_selector(
+                    module_tokens, tau=self.selector_tau, hard=not self.training,
+                )
+            if self.head_type == "module_mean":
+                cell_emb = module_tokens.mean(dim=1)
+            elif self.head_type == "max":
+                cell_emb = module_tokens.amax(dim=1)
+            elif self.head_type == "attn_pool":
+                cell_emb, _ = self.pooler(module_tokens)
+            elif self.head_type == "gene_mean":
+                gene_tokens = self._module_to_gene(module_tokens)
+                cell_emb = gene_tokens.mean(dim=1)
+            else:
+                raise ValueError(f"Unsupported head_type: {self.head_type}")
+
+        cell_emb = self.output_dropout(cell_emb)
+        return self.classifier(cell_emb)
+
+
+def get_encoder_blocks(backbone: "PerformerEncoderOnly_Module") -> List[nn.Module]:
+    net = backbone.encoder.net
+    if hasattr(net, "layers"):
+        return list(net.layers)
+    if hasattr(net, "blocks"):
+        return list(net.blocks)
+    raise AttributeError("Cannot find encoder blocks on Performer backbone")
+
+
+def set_requires_grad(module: nn.Module, flag: bool) -> None:
+    for param in module.parameters():
+        param.requires_grad = flag
+
+
+def configure_backbone_trainability(
+    backbone: "PerformerEncoderOnly_Module",
+    finetune_mode: str,
+    unfreeze_last_n: int,
+    train_pos_emb: bool,
+    train_module_embed: bool,
+) -> Dict[str, int]:
+    set_requires_grad(backbone, False)
+
+    if finetune_mode == "full":
+        # "full" here means full finetuning of the encoder pathway used by
+        # downstream classification, not the MAE-style reconstruction decoder.
+        set_requires_grad(backbone.module_embed, True)
+        set_requires_grad(backbone.encoder, True)
+        set_requires_grad(backbone.encoder_norm, True)
+        backbone.pos_emb.requires_grad = True
+        if getattr(backbone, "cls_token", None) is not None:
+            backbone.cls_token.requires_grad = True
+    elif finetune_mode == "last_n":
+        blocks = get_encoder_blocks(backbone)
+        if unfreeze_last_n <= 0:
+            raise ValueError("unfreeze_last_n must be > 0 when finetune_mode='last_n'")
+        for block in blocks[-min(unfreeze_last_n, len(blocks)):]:
+            set_requires_grad(block, True)
+        set_requires_grad(backbone.encoder_norm, True)
+    elif finetune_mode != "frozen":
+        raise ValueError(f"Unsupported finetune_mode: {finetune_mode}")
+
+    if train_pos_emb:
+        backbone.pos_emb.requires_grad = True
+        if getattr(backbone, "cls_token", None) is not None:
+            backbone.cls_token.requires_grad = True
+
+    if train_module_embed:
+        set_requires_grad(backbone.module_embed, True)
+
+    if getattr(backbone, 'family_emb', None) is not None:
+        backbone.family_emb.requires_grad = True
+    if getattr(backbone, 'module_gate', None) is not None:
+        set_requires_grad(backbone.module_gate, True)
+
+    total_params = sum(p.numel() for p in backbone.parameters())
+    trainable_params = sum(p.numel() for p in backbone.parameters() if p.requires_grad)
+    return {
+        "total": total_params,
+        "trainable": trainable_params,
+    }
+
+
+def build_optimizer(
+    model: PerformerModuleCellClassifier,
+    lr_head: float,
+    lr_backbone: float,
+    weight_decay: float,
+) -> torch.optim.Optimizer:
+    backbone_params = []
+    head_params = []
+    new_module_params = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if 'family_emb' in name or 'module_gate' in name or 'module_selector' in name:
+            new_module_params.append(param)
+        elif name.startswith("backbone."):
+            backbone_params.append(param)
+        else:
+            head_params.append(param)
+
+    param_groups = []
+    if backbone_params:
+        param_groups.append({"params": backbone_params, "lr": lr_backbone})
+    if head_params:
+        param_groups.append({"params": head_params, "lr": lr_head})
+    if new_module_params:
+        param_groups.append({"params": new_module_params, "lr": lr_head})
+
+    if not param_groups:
+        raise RuntimeError("No trainable parameters found")
+
+    return torch.optim.AdamW(param_groups, weight_decay=weight_decay, betas=(0.9, 0.999))
+
+
+def resolve_batch_sizes(args, world_size: int) -> Tuple[int, int, int]:
+    if args.global_batch_size is not None:
+        denom = world_size * args.grad_accum_steps
+        if args.global_batch_size % denom != 0:
+            raise ValueError(
+                f"global_batch_size={args.global_batch_size} must be divisible by "
+                f"world_size({world_size}) * grad_accum_steps({args.grad_accum_steps})"
+            )
+        train_batch_size = args.global_batch_size // denom
+    else:
+        train_batch_size = args.batch_size
+
+    if train_batch_size <= 0:
+        raise ValueError(f"Resolved per-device train batch size must be > 0, got {train_batch_size}")
+
+    if args.global_eval_batch_size is not None:
+        if args.global_eval_batch_size % world_size != 0:
+            raise ValueError(
+                f"global_eval_batch_size={args.global_eval_batch_size} must be divisible by world_size({world_size})"
+            )
+        eval_batch_size = args.global_eval_batch_size // world_size
+    elif args.eval_batch_size is not None:
+        eval_batch_size = args.eval_batch_size
+    else:
+        eval_batch_size = train_batch_size
+
+    if eval_batch_size <= 0:
+        raise ValueError(f"Resolved per-device eval batch size must be > 0, got {eval_batch_size}")
+
+    effective_global_batch = train_batch_size * world_size * args.grad_accum_steps
+    return train_batch_size, eval_batch_size, effective_global_batch
+
+
+def infer_model_config(args, checkpoint: dict, pkl_config: Optional[dict], module_data: List[List[int]]) -> dict:
+    from models_performer_module import sanitize_performer_module_config
+
+    if pkl_config is not None and isinstance(pkl_config, dict):
+        config = sanitize_performer_module_config(
+            pkl_config,
+            source=str(args.config_pkl),
+            log_fn=logger.info,
+        )
+        logger.info("Using config from PKL: %s", args.config_pkl)
+        if "module_sets" in config:
+            module_data = config["module_sets"]
+    elif isinstance(checkpoint, dict) and "config" in checkpoint:
+        config = sanitize_performer_module_config(
+            checkpoint["config"],
+            source=str(resolve_path(args.performer_checkpoint)),
+            log_fn=logger.info,
+        )
+        logger.info("Using config from checkpoint")
+    else:
+        raise ValueError(
+            "Backbone config was not found in config_pkl or checkpoint['config']. "
+            "This script now expects pretrained structural args to be restored automatically."
+        )
+
+    return {
+        "module_sets": module_data,
+        "num_genes": config["num_genes"],
+        "base_dim": config["base_dim"],
+        "enc_depth": config["enc_depth"],
+        "enc_heads": config["enc_heads"],
+        "enc_dim_head": config["enc_dim_head"],
+        "gene_embed_dim": config.get("gene_embed_dim", 200),
+        "fusion_method": config.get("fusion_method", "multiply"),
+        "use_gene2vec": config.get("use_gene2vec", True),
+        "loss_mode": config.get("loss_mode", "gene_mean"),
+        "recon_target_level": config.get("recon_target_level", "gene"),
+        "pretrain_arch": config.get("pretrain_arch", "mae"),
+        "mask_style": config.get("mask_style", "drop"),
+        "mask_token_mode": config.get("mask_token_mode", "learned"),
+        "pos_embed_mode": config.get("pos_embed_mode", "module_id"),
+        "decoder_type": config.get("decoder_type", "simple"),
+        "decoder_gene_agg": config.get("decoder_gene_agg", False),
+        "dec_embed_dim": config.get("dec_embed_dim", config["base_dim"]),
+        "dec_depth": config.get("dec_depth", 4),
+        "dec_heads": config.get("dec_heads", 8),
+        "dec_dim_head": config.get("dec_dim_head", 64),
+        "use_cls_token": bool(config.get("use_cls_token", False)),
+        "cls_loss_weight": float(config.get("cls_loss_weight", 0.0)),
+        "cls_head_hidden_dim": int(config.get("cls_head_hidden_dim", 0)),
+    }
+
+
+def compute_module_family_ids(module_data_dir: str) -> List[int]:
+    """Read merged_modules.json, return family ID (0-3) for each module."""
+    merged_path = Path(module_data_dir) / "merged_modules.json"
+    with open(merged_path) as f:
+        modules = json.load(f)
+    PREFIXES = [
+        ("module1_hallmarks", 0),
+        ("module2_reactome", 1),
+        ("module3_hotspot", 2),
+        ("module3_processed", 3),
+    ]
+    family_ids = []
+    for name in modules.keys():
+        fid = 0
+        for prefix, idx in PREFIXES:
+            if name.startswith(prefix):
+                fid = idx
+                break
+        family_ids.append(fid)
+    return family_ids
+
+
+def load_backbone(args, device: torch.device) -> Tuple["PerformerEncoderOnly_Module", List[List[int]], dict]:
+    from models_performer_module import PerformerEncoderOnly_Module
+    from module_dataset import load_module_data
+
+    config_pkl = None
+    config_path = resolve_path(args.config_pkl)
+    if config_path.exists():
+        with open(config_path, "rb") as handle:
+            config_pkl = migrate_legacy_module_config(pickle.load(handle))
+    else:
+        logger.warning("Config PKL not found at %s; fallback to checkpoint config", config_path)
+
+    ckpt_path = resolve_path(args.performer_checkpoint)
+    checkpoint = torch.load(ckpt_path, map_location=device)
+
+    module_data = None
+    if isinstance(config_pkl, dict) and "module_sets" in config_pkl:
+        module_data = config_pkl["module_sets"]
+        logger.info("Using module_sets from PKL (%d modules)", len(module_data))
+    elif isinstance(checkpoint, dict) and isinstance(checkpoint.get("config"), dict) and "module_sets" in checkpoint["config"]:
+        module_data = checkpoint["config"]["module_sets"]
+        logger.info("Using module_sets from checkpoint config (%d modules)", len(module_data))
+    else:
+        if not args.module_data_dir:
+            raise ValueError("module_data_dir is required when module_sets are not stored in the PKL/checkpoint config")
+        module_data = load_module_data(
+            module_data_dir=str(resolve_path(args.module_data_dir)),
+            gene_list_path=str(resolve_path(args.gene_list_path)),
+            gene_mapping_path=str(resolve_path(args.gene_mapping_path)),
+            verbose=True,
+        )
+
+    if args.shuffle_module_genes:
+        logger.info("Shuffling genes across modules with seed=%d", args.shuffle_seed)
+        rng = random.Random(args.shuffle_seed)
+        module_sizes = [len(module) for module in module_data]
+        all_genes = [gene for module in module_data for gene in module]
+        rng.shuffle(all_genes)
+        shuffled = []
+        start = 0
+        for size in module_sizes:
+            shuffled.append(all_genes[start:start + size])
+            start += size
+        module_data = shuffled
+
+    model_config = infer_model_config(args, checkpoint, config_pkl, module_data)
+    module_data = model_config["module_sets"]
+
+    if "model_state_dict" in checkpoint:
+        state_dict = checkpoint["model_state_dict"]
+    elif "model" in checkpoint:
+        state_dict = checkpoint["model"]
+    else:
+        state_dict = checkpoint
+
+    has_cls_in_checkpoint = "cls_token" in state_dict
+    use_cls_token = has_cls_in_checkpoint or model_config["use_cls_token"]
+
+    logger.info(
+        "Backbone config: base_dim=%d depth=%d heads=%d modules=%d cls=%s",
+        model_config["base_dim"],
+        model_config["enc_depth"],
+        model_config["enc_heads"],
+        len(module_data),
+        use_cls_token,
+    )
+
+    module_family_ids = None
+    if getattr(args, "use_family_emb", False):
+        pdd = getattr(args, "module_data_dir", None)
+        if pdd:
+            module_family_ids = compute_module_family_ids(str(resolve_path(pdd)))
+            logger.info("Family embedding enabled: %d module family IDs computed", len(module_family_ids))
+
+    backbone = PerformerEncoderOnly_Module(
+        module_sets=module_data,
+        num_genes=model_config["num_genes"],
+        max_seq_len=len(module_data),
+        base_dim=model_config["base_dim"],
+        enc_depth=model_config["enc_depth"],
+        enc_heads=model_config["enc_heads"],
+        enc_dim_head=model_config["enc_dim_head"],
+        mask_ratio=0.75,
+        gene2vec_file=str(resolve_path(args.gene2vec_data_path)),
+        gene_embed_dim=model_config["gene_embed_dim"],
+        fusion_method=model_config["fusion_method"],
+        use_cls_token=use_cls_token,
+        cls_loss_weight=model_config["cls_loss_weight"] if use_cls_token else 0.0,
+        cls_head_hidden_dim=model_config["cls_head_hidden_dim"],
+        use_gene2vec=model_config["use_gene2vec"],
+        loss_mode=model_config["loss_mode"],
+        recon_target_level=model_config["recon_target_level"],
+        pretrain_arch=model_config["pretrain_arch"],
+        mask_style=model_config["mask_style"],
+        mask_token_mode=model_config["mask_token_mode"],
+        pos_embed_mode=model_config["pos_embed_mode"],
+        decoder_type=model_config["decoder_type"],
+        decoder_gene_agg=model_config["decoder_gene_agg"],
+        dec_embed_dim=model_config["dec_embed_dim"],
+        dec_depth=model_config["dec_depth"],
+        dec_heads=model_config["dec_heads"],
+        dec_dim_head=model_config["dec_dim_head"],
+        use_family_emb=getattr(args, "use_family_emb", False),
+        module_family_ids=module_family_ids,
+        use_module_gate=getattr(args, "use_module_gate", False),
+    )
+
+    filtered_state_dict = {
+        key: value
+        for key, value in state_dict.items()
+        if "go_conv" not in key and "G_go" not in key.lower()
+    }
+    missing_keys, unexpected_keys = backbone.load_state_dict(filtered_state_dict, strict=False)
+    expected_new = {"family_emb", "module_family_ids", "module_gate."}
+    real_missing = [k for k in missing_keys
+                    if not any(k == e or k.startswith(e) for e in expected_new)]
+    if real_missing:
+        logger.warning("Unexpected missing keys: %s", real_missing[:10])
+    if unexpected_keys:
+        logger.warning("Unexpected keys: %s", unexpected_keys[:10])
+
+    backbone.to(device)
+    return backbone, module_data, model_config
+
+
+def stratified_split(
+    labels: np.ndarray,
+    seed: int,
+    val_ratio: float,
+    test_ratio: float,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    indices = np.arange(len(labels))
+    test_split = StratifiedShuffleSplit(n_splits=1, test_size=test_ratio, random_state=seed)
+    train_val_idx, test_idx = next(test_split.split(indices, labels))
+
+    remaining_ratio = 1.0 - test_ratio
+    if remaining_ratio <= 0.0:
+        raise ValueError("test_ratio must be < 1")
+    val_ratio_adjusted = val_ratio / remaining_ratio
+
+    train_val_labels = labels[train_val_idx]
+    val_split = StratifiedShuffleSplit(n_splits=1, test_size=val_ratio_adjusted, random_state=seed)
+    train_sub_idx, val_sub_idx = next(val_split.split(train_val_idx, train_val_labels))
+
+    train_idx = train_val_idx[train_sub_idx]
+    val_idx = train_val_idx[val_sub_idx]
+    return train_idx, val_idx, test_idx
+
+
+def stratified_train_val_split(
+    labels: np.ndarray,
+    seed: int,
+    val_ratio: float,
+) -> Tuple[np.ndarray, np.ndarray]:
+    if not 0.0 < val_ratio < 1.0:
+        raise ValueError(f"val_ratio must be in (0, 1), got {val_ratio}")
+    indices = np.arange(len(labels))
+    val_split = StratifiedShuffleSplit(n_splits=1, test_size=val_ratio, random_state=seed)
+    train_idx, val_idx = next(val_split.split(indices, labels))
+    return train_idx, val_idx
+
+
+@torch.no_grad()
+def evaluate(
+    model: PerformerModuleCellClassifier,
+    data_loader: DataLoader,
+    preprocessor: DataPreprocessor,
+    reindex_state: dict,
+    criterion: nn.Module,
+    device: torch.device,
+) -> dict:
+    model.eval()
+    loss_sum = 0.0
+    sample_count = 0
+    preds_all = []
+    labels_all = []
+
+    for x, y in data_loader:
+        x = x.to(device)
+        y = y.to(device)
+        x_model = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
+        logits = model(x_model)
+        loss = criterion(logits, y)
+        loss_sum += loss.item() * y.size(0)
+        sample_count += y.size(0)
+        preds_all.append(logits.argmax(dim=-1).cpu().numpy())
+        labels_all.append(y.cpu().numpy())
+
+    preds = np.concatenate(preds_all) if preds_all else np.empty((0,), dtype=np.int64)
+    labels = np.concatenate(labels_all) if labels_all else np.empty((0,), dtype=np.int64)
+    preds = all_gather_numpy(preds)
+    labels = all_gather_numpy(labels)
+    loss_sum = reduce_sum(loss_sum, device)
+    sample_count = int(reduce_sum(sample_count, device))
+    return {
+        "loss": float(loss_sum / max(sample_count, 1)),
+        "acc": float(accuracy_score(labels, preds)),
+        "macro_f1": float(f1_score(labels, preds, average="macro")),
+        "preds": preds,
+        "labels": labels,
+    }
+
+
+def train_one_epoch(
+    model: PerformerModuleCellClassifier,
+    data_loader: DataLoader,
+    preprocessor: DataPreprocessor,
+    reindex_state: dict,
+    criterion: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    device: torch.device,
+    grad_clip: float,
+    grad_accum_steps: int,
+    use_ddp: bool,
+    selector_sparsity_weight: float = 0.0,
+) -> dict:
+    model.train()
+    running_loss = 0.0
+    sample_count = 0
+    preds_all = []
+    labels_all = []
+    optimizer.zero_grad(set_to_none=True)
+
+    num_steps = len(data_loader)
+    for step, (x, y) in enumerate(data_loader, start=1):
+        x = x.to(device)
+        y = y.to(device)
+        x_model = preprocessor.reindex_tensor_to_model_space(x, reindex_state=reindex_state)
+
+        sync_context = (
+            model.no_sync()
+            if use_ddp and grad_accum_steps > 1 and step % grad_accum_steps != 0
+            else contextlib.nullcontext()
+        )
+        with sync_context:
+            logits = model(x_model)
+            loss = criterion(logits, y)
+            if selector_sparsity_weight > 0:
+                model_raw = model.module if use_ddp else model
+                loss = loss + selector_sparsity_weight * model_raw.get_selector_sparsity_loss()
+            (loss / grad_accum_steps).backward()
+
+        if step % grad_accum_steps == 0 or step == num_steps:
+            if grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+
+        running_loss += loss.item() * y.size(0)
+        sample_count += y.size(0)
+        preds_all.append(logits.argmax(dim=-1).detach().cpu().numpy())
+        labels_all.append(y.detach().cpu().numpy())
+
+    preds = np.concatenate(preds_all) if preds_all else np.empty((0,), dtype=np.int64)
+    labels = np.concatenate(labels_all) if labels_all else np.empty((0,), dtype=np.int64)
+    preds = all_gather_numpy(preds)
+    labels = all_gather_numpy(labels)
+    running_loss = reduce_sum(running_loss, device)
+    sample_count = int(reduce_sum(sample_count, device))
+    return {
+        "loss": float(running_loss / max(sample_count, 1)),
+        "acc": float(accuracy_score(labels, preds)),
+        "macro_f1": float(f1_score(labels, preds, average="macro")),
+    }
+
+
+def save_json(path: Path, payload: dict) -> None:
+    with open(path, "w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+
+
+def build_confusion_tables(
+    labels: np.ndarray,
+    preds: np.ndarray,
+    class_names: Sequence[str],
+) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame, pd.DataFrame]:
+    label_ids = np.arange(len(class_names), dtype=np.int64)
+    conf = confusion_matrix(labels, preds, labels=label_ids)
+    row_sums = conf.sum(axis=1, keepdims=True)
+    conf_norm = np.divide(
+        conf.astype(np.float64),
+        row_sums,
+        out=np.zeros_like(conf, dtype=np.float64),
+        where=row_sums != 0,
+    )
+    conf_df = pd.DataFrame(conf, index=class_names, columns=class_names)
+    conf_norm_df = pd.DataFrame(conf_norm, index=class_names, columns=class_names)
+    return conf, conf_norm, conf_df, conf_norm_df
+
+
+def apply_dataset_preset(args) -> None:
+    if args.dataset_preset == "none":
+        return
+
+    if args.dataset_preset == "ms":
+        default_single_data_path = "../../cell_annotation/cellxgene80k_balanced_10types.h5ad"
+        default_train_path = "../../cell_annotation/data/c_data.h5ad"
+        default_test_path = "../../cell_annotation/data/filtered_ms_adata.h5ad"
+
+        if args.data_path == default_single_data_path and args.train_data_path is None and args.test_data_path is None:
+            args.train_data_path = default_train_path
+            args.test_data_path = default_test_path
+
+        if args.label_column == "cell_type_anno":
+            args.label_column = "celltype"
+        if args.test_label_column is None:
+            args.test_label_column = "celltype"
+        if args.gene_name_column == "auto":
+            args.gene_name_column = "gene_name"
+        if args.min_genes == 10:
+            args.min_genes = 0
+        if args.min_cells == 10:
+            args.min_cells = 0
+
+        logger.info(
+            "Applied dataset preset 'ms': train_data_path=%s test_data_path=%s label_column=%s "
+            "test_label_column=%s gene_name_column=%s min_genes=%d min_cells=%d",
+            args.train_data_path,
+            args.test_data_path,
+            args.label_column,
+            args.test_label_column,
+            args.gene_name_column,
+            args.min_genes,
+            args.min_cells,
+        )
+        return
+
+    raise ValueError(f"Unsupported dataset_preset: {args.dataset_preset}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="End-to-end fine-tuning for Performer Module cell annotation")
+
+    parser.add_argument("--seed", type=int, default=2021)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--batch_size", type=int, default=32, help="Per-device train batch size.")
+    parser.add_argument("--eval_batch_size", type=int, default=None, help="Per-device eval batch size.")
+    parser.add_argument("--global_batch_size", type=int, default=None, help="Effective global train batch size.")
+    parser.add_argument("--global_eval_batch_size", type=int, default=None, help="Effective global eval batch size.")
+    parser.add_argument("--grad_accum_steps", type=int, default=1)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--prefetch_factor", type=int, default=2)
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--patience", type=int, default=10)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+
+    parser.add_argument("--lr_head", type=float, default=1e-4)
+    parser.add_argument("--lr_backbone", type=float, default=1e-5)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--label_smoothing", type=float, default=0.1)
+    parser.add_argument(
+        "--dropout",
+        type=float,
+        default=0.2,
+        help="Shared dropout applied to pooled cell embeddings before classification.",
+    )
+    parser.add_argument("--pool_hidden_dim", type=int, default=0)
+    parser.add_argument("--cls_decoder_layers", type=int, default=3)
+
+    parser.add_argument(
+        "--head_type",
+        type=str,
+        default="attn_pool",
+        choices=["module_mean", "max", "gene_mean", "attn_pool", "cls"],
+    )
+    parser.add_argument(
+        "--finetune_mode",
+        type=str,
+        default="last_n",
+        choices=["frozen", "last_n", "full"],
+    )
+    parser.add_argument("--unfreeze_last_n", type=int, default=1)
+    parser.add_argument("--train_pos_emb", action="store_true")
+    parser.add_argument("--train_module_embed", action="store_true")
+    parser.add_argument("--use_family_emb", action="store_true",
+                        help="Add family-type embedding (Hallmarks/Reactome/Hotspot/WGCNA)")
+    parser.add_argument("--use_module_gate", action="store_true",
+                        help="Add learnable sparse gate per module")
+    parser.add_argument("--use_module_selector", action="store_true",
+                        help="Dynamic module selector with sparsity loss (soft train, hard inference)")
+    parser.add_argument("--selector_sparsity_weight", type=float, default=0.1,
+                        help="Lambda for sparsity loss on module selector scores")
+    parser.add_argument("--selector_tau", type=float, default=0.5,
+                        help="Score threshold for hard module selection during inference")
+    parser.add_argument("--dataset_preset", type=str, default="none", choices=["none", "ms"])
+
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        default="../../cell_annotation/cellxgene80k_balanced_10types.h5ad",
+    )
+    parser.add_argument("--train_data_path", type=str, default=None)
+    parser.add_argument("--test_data_path", type=str, default=None)
+    parser.add_argument("--label_column", type=str, default="cell_type_anno")
+    parser.add_argument("--test_label_column", type=str, default=None)
+    parser.add_argument(
+        "--gene_name_column",
+        type=str,
+        default="auto",
+        help="Gene-name source in adata.var: 'auto', 'index', or a column name such as 'gene_name'.",
+    )
+    parser.add_argument("--val_ratio", type=float, default=0.1)
+    parser.add_argument("--test_ratio", type=float, default=0.2)
+    parser.add_argument("--min_genes", type=int, default=10)
+    parser.add_argument("--min_cells", type=int, default=10)
+    parser.add_argument("--max_cells", type=int, default=None)
+
+    parser.add_argument("--config_pkl", type=str, default="../ckpts_performer_module_test/performer_module_test.pkl")
+    parser.add_argument(
+        "--performer_checkpoint",
+        type=str,
+        default="../ckpts_performer_module_test/performer_module_test_epoch7.pth",
+    )
+    parser.add_argument("--module_data_dir", type=str, default="../module_data_all")
+    parser.add_argument("--gene_list_path", type=str, default="../../selected_genes_27k.txt")
+    parser.add_argument("--gene_mapping_path", type=str, default="../../gene2ens.csv")
+    parser.add_argument("--gene2vec_data_path", type=str, default="../../selected_gene2vec_27k.npy")
+    parser.add_argument("--gene2ens_file", type=str, default="../../gene2ens.csv")
+    parser.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
+    parser.add_argument(
+        "--shuffle_module_genes",
+        action="store_true",
+        help="Apply an extra downstream reshuffle on top of module_sets restored from PKL. "
+             "Do not pass this when reproducing the pretrained backbone as-is.",
+    )
+    parser.add_argument(
+        "--shuffle_seed",
+        type=int,
+        default=2021,
+        help="Random seed for the downstream extra reshuffle when --shuffle_module_genes is enabled.",
+    )
+
+    parser.add_argument("--output_dir", type=str, default="./ckpts_performer_module_finetune")
+    parser.add_argument("--run_name", type=str, default="")
+    parser.add_argument(
+        "--save_test_artifacts",
+        action="store_true",
+        help="Save test classification report, confusion matrices, and metrics summary.",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    apply_dataset_preset(args)
+    device, use_ddp, rank, world_size, local_rank = init_distributed_mode(args.device)
+    if not is_main_process():
+        logger.setLevel(logging.WARNING)
+    seed_all(args.seed + rank)
+
+    run_name = args.run_name or (
+        f"{args.head_type}_{args.finetune_mode}"
+        f"_n{args.unfreeze_last_n}_lrh{args.lr_head:g}_lrb{args.lr_backbone:g}_seed{args.seed}"
+    )
+    output_dir = resolve_path(args.output_dir) / run_name
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    train_batch_size, eval_batch_size, effective_global_batch = resolve_batch_sizes(args, world_size)
+    logger.info("Run directory: %s", output_dir)
+    logger.info(
+        "Device: %s | DDP=%s rank=%d world_size=%d local_rank=%d",
+        device,
+        use_ddp,
+        rank,
+        world_size,
+        local_rank,
+    )
+    logger.info(
+        "Batch config: per_device_train=%d per_device_eval=%d grad_accum=%d effective_global_train=%d",
+        train_batch_size,
+        eval_batch_size,
+        args.grad_accum_steps,
+        effective_global_batch,
+    )
+
+    preprocessor = DataPreprocessor(
+        gene2ens_file=str(resolve_path(args.gene2ens_file)),
+        gene2ens_column=args.gene2ens_column,
+        model_genes_file=str(resolve_path(args.gene_list_path)),
+    )
+    preprocessor.load_gene_mapping()
+    preprocessor.load_model_genes()
+    paired_data_mode = bool(args.train_data_path or args.test_data_path)
+    if bool(args.train_data_path) != bool(args.test_data_path):
+        raise ValueError("train_data_path and test_data_path must be provided together")
+
+    if paired_data_mode:
+        if args.test_ratio != 0.2:
+            logger.info("External test set provided; ignoring test_ratio=%s", args.test_ratio)
+
+        train_adata, train_reindex_state = preprocessor.preprocess_adata(
+            str(resolve_path(args.train_data_path)),
+            min_genes=args.min_genes,
+            min_cells=args.min_cells,
+            max_cells=args.max_cells,
+            dataset_name="train",
+            gene_name_column=args.gene_name_column,
+        )
+        test_adata, test_reindex_state = preprocessor.preprocess_adata(
+            str(resolve_path(args.test_data_path)),
+            min_genes=args.min_genes,
+            min_cells=args.min_cells,
+            max_cells=None,
+            dataset_name="test",
+            gene_name_column=args.gene_name_column,
+        )
+
+        train_label_column = args.label_column
+        test_label_column = args.test_label_column or args.label_column
+        label_encoder = LabelEncoder()
+        train_labels_raw = train_adata.obs[train_label_column].astype(str).to_numpy()
+        label_encoder.fit(train_labels_raw)
+        labels_train = label_encoder.transform(train_labels_raw)
+
+        test_labels_raw = test_adata.obs[test_label_column].astype(str).to_numpy()
+        unseen_test_labels = sorted(set(test_labels_raw) - set(label_encoder.classes_))
+        if unseen_test_labels:
+            raise ValueError(
+                "External test set contains labels unseen in training data: "
+                + ", ".join(unseen_test_labels[:10])
+            )
+        labels_test = label_encoder.transform(test_labels_raw)
+        num_classes = len(label_encoder.classes_)
+
+        train_idx, val_idx = stratified_train_val_split(
+            labels=labels_train,
+            seed=args.seed,
+            val_ratio=args.val_ratio,
+        )
+        logger.info(
+            "Using external test split | train=%s test=%s",
+            resolve_path(args.train_data_path),
+            resolve_path(args.test_data_path),
+        )
+        logger.info(
+            "Split sizes: train=%d val=%d external_test=%d classes=%d",
+            len(train_idx),
+            len(val_idx),
+            test_adata.n_obs,
+            num_classes,
+        )
+
+        train_dataset = ExpressionDataset(train_adata.X, labels_train, train_idx)
+        val_dataset = ExpressionDataset(train_adata.X, labels_train, val_idx)
+        test_dataset = ExpressionDataset(test_adata.X, labels_test, np.arange(test_adata.n_obs))
+        val_reindex_state = train_reindex_state
+    else:
+        adata, train_reindex_state = preprocessor.preprocess_adata(
+            str(resolve_path(args.data_path)),
+            min_genes=args.min_genes,
+            min_cells=args.min_cells,
+            max_cells=args.max_cells,
+            dataset_name="data",
+            gene_name_column=args.gene_name_column,
+        )
+
+        label_encoder = LabelEncoder()
+        labels = label_encoder.fit_transform(adata.obs[args.label_column].astype(str).to_numpy())
+        num_classes = len(label_encoder.classes_)
+        train_idx, val_idx, test_idx = stratified_split(
+            labels=labels,
+            seed=args.seed,
+            val_ratio=args.val_ratio,
+            test_ratio=args.test_ratio,
+        )
+        logger.info(
+            "Split sizes: train=%d val=%d test=%d classes=%d",
+            len(train_idx),
+            len(val_idx),
+            len(test_idx),
+            num_classes,
+        )
+
+        matrix = adata.X
+        train_dataset = ExpressionDataset(matrix, labels, train_idx)
+        val_dataset = ExpressionDataset(matrix, labels, val_idx)
+        test_dataset = ExpressionDataset(matrix, labels, test_idx)
+        val_reindex_state = train_reindex_state
+        test_reindex_state = train_reindex_state
+
+    train_sampler = (
+        DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False)
+        if use_ddp
+        else None
+    )
+    val_sampler = (
+        DistributedEvalSampler(len(val_dataset), rank=rank, world_size=world_size)
+        if use_ddp
+        else None
+    )
+    test_sampler = (
+        DistributedEvalSampler(len(test_dataset), rank=rank, world_size=world_size)
+        if use_ddp
+        else None
+    )
+
+    loader_kwargs = {
+        "num_workers": args.num_workers,
+        "pin_memory": device.type == "cuda",
+    }
+    if args.num_workers > 0:
+        loader_kwargs["persistent_workers"] = True
+        loader_kwargs["prefetch_factor"] = args.prefetch_factor
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=train_batch_size,
+        shuffle=train_sampler is None,
+        sampler=train_sampler,
+        **loader_kwargs,
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=eval_batch_size,
+        shuffle=False,
+        sampler=val_sampler,
+        **loader_kwargs,
+    )
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=eval_batch_size,
+        shuffle=False,
+        sampler=test_sampler,
+        **loader_kwargs,
+    )
+
+    backbone, module_sets, backbone_config = load_backbone(args, device=device)
+    trainability = configure_backbone_trainability(
+        backbone,
+        finetune_mode=args.finetune_mode,
+        unfreeze_last_n=args.unfreeze_last_n,
+        train_pos_emb=args.train_pos_emb,
+        train_module_embed=args.train_module_embed,
+    )
+
+    model = PerformerModuleCellClassifier(
+        backbone=backbone,
+        module_sets=module_sets,
+        num_classes=num_classes,
+        head_type=args.head_type,
+        pool_hidden_dim=args.pool_hidden_dim,
+        dropout=args.dropout,
+        cls_decoder_layers=args.cls_decoder_layers,
+        use_module_selector=getattr(args, "use_module_selector", False),
+    ).to(device)
+    model_to_save = model
+    if getattr(args, "use_module_selector", False):
+        model_to_save.selector_tau = getattr(args, "selector_tau", 0.5)
+
+    head_total = sum(p.numel() for name, p in model.named_parameters() if not name.startswith("backbone."))
+    head_trainable = sum(
+        p.numel() for name, p in model.named_parameters() if not name.startswith("backbone.") and p.requires_grad
+    )
+    logger.info(
+        "Backbone params: trainable=%d total=%d | Head params: trainable=%d total=%d",
+        trainability["trainable"],
+        trainability["total"],
+        head_trainable,
+        head_total,
+    )
+
+    optimizer = build_optimizer(
+        model=model,
+        lr_head=args.lr_head,
+        lr_backbone=args.lr_backbone,
+        weight_decay=args.weight_decay,
+    )
+    if use_ddp:
+        model = DDP(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(args.epochs, 1))
+    criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing)
+
+    if is_main_process():
+        save_json(
+            output_dir / "run_config.json",
+            {
+                "args": vars(args),
+                "run_name": run_name,
+                "num_classes": num_classes,
+                "label_classes": label_encoder.classes_.tolist(),
+                "backbone_config": backbone_config,
+                "trainable_backbone_params": trainability,
+                "head_type": args.head_type,
+                "data_mode": "paired_train_test" if paired_data_mode else "single_split",
+                "gene_name_column": args.gene_name_column,
+                "world_size": world_size,
+                "effective_global_batch": effective_global_batch,
+                "per_device_batch_size": train_batch_size,
+                "per_device_eval_batch_size": eval_batch_size,
+            },
+        )
+    if use_ddp:
+        dist.barrier()
+
+    best_val_f1 = -1.0
+    best_metrics = None
+    best_epoch = -1
+    patience_counter = 0
+
+    for epoch in range(1, args.epochs + 1):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        train_metrics = train_one_epoch(
+            model=model,
+            data_loader=train_loader,
+            preprocessor=preprocessor,
+            reindex_state=train_reindex_state,
+            criterion=criterion,
+            optimizer=optimizer,
+            device=device,
+            grad_clip=args.grad_clip,
+            grad_accum_steps=args.grad_accum_steps,
+            use_ddp=use_ddp,
+            selector_sparsity_weight=getattr(args, "selector_sparsity_weight", 0.0) if getattr(args, "use_module_selector", False) else 0.0,
+        )
+        val_metrics = evaluate(
+            model=model,
+            data_loader=val_loader,
+            preprocessor=preprocessor,
+            reindex_state=val_reindex_state,
+            criterion=criterion,
+            device=device,
+        )
+        scheduler.step()
+
+        logger.info(
+            "Epoch %d | train loss %.4f acc %.4f f1 %.4f | val loss %.4f acc %.4f f1 %.4f",
+            epoch,
+            train_metrics["loss"],
+            train_metrics["acc"],
+            train_metrics["macro_f1"],
+            val_metrics["loss"],
+            val_metrics["acc"],
+            val_metrics["macro_f1"],
+        )
+
+        if val_metrics["macro_f1"] > best_val_f1:
+            best_val_f1 = val_metrics["macro_f1"]
+            best_epoch = epoch
+            patience_counter = 0
+
+            best_metrics = {
+                "epoch": epoch,
+                "train": train_metrics,
+                "val": {k: v for k, v in val_metrics.items() if k not in {"preds", "labels"}},
+            }
+            if is_main_process():
+                torch.save(
+                    {
+                        "epoch": epoch,
+                        "model_state_dict": unwrap_model(model_to_save).state_dict(),
+                        "optimizer_state_dict": optimizer.state_dict(),
+                        "scheduler_state_dict": scheduler.state_dict(),
+                        "best_val_f1": best_val_f1,
+                        "args": vars(args),
+                        "label_classes": label_encoder.classes_.tolist(),
+                        "backbone_config": backbone_config,
+                        "world_size": world_size,
+                        "effective_global_batch": effective_global_batch,
+                    },
+                    output_dir / "best_model.pt",
+                )
+        else:
+            patience_counter += 1
+            if patience_counter >= args.patience:
+                logger.info("Early stopping at epoch %d", epoch)
+                break
+        if use_ddp:
+            dist.barrier()
+
+    if best_epoch < 0:
+        raise RuntimeError("Training finished without a valid checkpoint")
+
+    if use_ddp:
+        dist.barrier()
+    ckpt = torch.load(output_dir / "best_model.pt", map_location=device)
+    unwrap_model(model).load_state_dict(ckpt["model_state_dict"], strict=True)
+
+    test_metrics = evaluate(
+        model=model,
+        data_loader=test_loader,
+        preprocessor=preprocessor,
+        reindex_state=test_reindex_state,
+        criterion=criterion,
+        device=device,
+    )
+    logger.info(
+        "Best epoch %d | test loss %.4f acc %.4f f1 %.4f",
+        best_epoch,
+        test_metrics["loss"],
+        test_metrics["acc"],
+        test_metrics["macro_f1"],
+    )
+
+    report = classification_report(
+        test_metrics["labels"],
+        test_metrics["preds"],
+        target_names=label_encoder.classes_.tolist(),
+        digits=4,
+    )
+    conf, conf_norm, conf_df, conf_norm_df = build_confusion_tables(
+        test_metrics["labels"],
+        test_metrics["preds"],
+        label_encoder.classes_.tolist(),
+    )
+
+    summary = {
+        "best_epoch": best_epoch,
+        "best_val_f1": best_val_f1,
+        "train": best_metrics["train"] if best_metrics is not None else None,
+        "val": best_metrics["val"] if best_metrics is not None else None,
+        "test": {k: v for k, v in test_metrics.items() if k not in {"preds", "labels"}},
+    }
+    if is_main_process():
+        logger.info("Test classification report:\n%s", report)
+        logger.info(
+            "Test confusion matrix (rows=true, cols=pred):\n%s",
+            conf_df.to_string(),
+        )
+        logger.info(
+            "Test confusion matrix normalized by true class:\n%s",
+            conf_norm_df.round(4).to_string(),
+        )
+        if args.save_test_artifacts:
+            with open(output_dir / "test_classification_report.txt", "w", encoding="utf-8") as handle:
+                handle.write(report)
+            with open(output_dir / "test_confusion_matrix.txt", "w", encoding="utf-8") as handle:
+                handle.write(conf_df.to_string())
+                handle.write("\n")
+            with open(output_dir / "test_confusion_matrix_normalized.txt", "w", encoding="utf-8") as handle:
+                handle.write(conf_norm_df.round(4).to_string())
+                handle.write("\n")
+            np.save(output_dir / "test_confusion_matrix.npy", conf)
+            np.save(output_dir / "test_confusion_matrix_normalized.npy", conf_norm)
+            conf_df.to_csv(output_dir / "test_confusion_matrix.csv", encoding="utf-8")
+            conf_norm_df.to_csv(output_dir / "test_confusion_matrix_normalized.csv", encoding="utf-8")
+            save_json(output_dir / "metrics_summary.json", summary)
+            logger.info("Saved test artifacts to %s", output_dir)
+        else:
+            logger.info("Test artifacts were not saved. Use --save_test_artifacts to enable file output.")
+    if use_ddp:
+        dist.barrier()
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

BioMod/command.txt

@@ -0,0 +1,31 @@
+source activate base
+conda activate sclong
+cd ../data/250010166/scLong/performer_module
+torchrun --nproc_per_node $SENSECORE_ACCELERATE_DEVICE_COUNT --nnodes $SENSECORE_PYTORCH_NNODES --node_rank $SENSECORE_PYTORCH_NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
+performer_module.py \
+    --gene_list_path ../selected_genes_27k.txt \
+    --cell_num 4000000 \
+    --epoch 15 \
+    --batch_size 60 \
+    --valid_batch_size 120 \
+    --learning_rate 1e-4 \
+    --grad_acc 5 \
+    --log_every 100 \
+    --mask_ratio 0.75 \
+    --data_path ../../../share/cell_files \
+    --module_data_dir module_data_all/ \
+    --gene_mapping_path ../gene2ens.csv \
+    --gene2vec_data_path ../selected_gene2vec_27k.npy \
+    --ckpt_dir ./performer_module_4M_all_module_hvg_600_singleton/ \
+    --model_name performer_module_4M_all_module_hvg_600_singleton \
+    --fusion_method multiply \
+    -use_wandb \
+    --api_key "$WANDB_API_KEY" \
+    --loss_mode gene_mean \
+    --decoder_type simple \
+    --num_workers 12 \
+    --prefetch_factor 4 \
+    --extra_gene_rank_path hvg_rankings/hvg_seurat_v3_top4096_first4000000cells_gene27874.npz \
+    --extra_gene_num 600 \
+    --extra_gene_module_strategy singleton \
+    --dec_embed_dim 1024

BioMod/drug_combination/README.md

@@ -0,0 +1,130 @@
+# Drug Combination Synergy Prediction (PerformerModule)
+
+This task predicts drug-combination synergy using pretrained PerformerModule cell
+embeddings and GCN-based drug molecular graph encoders.
+
+It is adapted from [scLong/drug_combination](../../../drug_combination)
+(DeepDDs), replacing the scfm/geneformer cell-line representation with
+PerformerModule.
+
+## Model Architecture
+
+| Component | Description |
+| --- | --- |
+| Drug branch | Three PyG GCNConv layers (78 -> 78 -> 156 -> 312), `global_max_pool`, and two fully connected layers to 128 dimensions |
+| Cell branch | Pretrained PerformerModule encoder, module mean pooling, and an MLP to 128 dimensions |
+| Fusion layer | `concat(drug1, drug2, cell) = 384 -> MLP -> 2-class prediction` |
+
+The drug1 and drug2 branches share weights, matching scLong.
+
+## Data
+
+Place the data under `data_summary/`:
+
+```text
+data_summary/
+|-- new_labels_0_10.csv          # Training labels (drug1, drug2, cell, label)
+|-- new_cell_features_954.csv    # Training cell features (30 cell lines x 954 genes)
+|-- smiles_0_10.csv              # Drug SMILES
+`-- test_data/
+    |-- independent_input.csv              # Independent test labels (668 samples)
+    `-- independent_cell_features_954.csv  # Test cell features (1020 cell lines x 954 genes)
+```
+
+## Dependencies
+
+- PyTorch
+- PyTorch Geometric (`torch_geometric`)
+- RDKit
+- networkx
+- scikit-learn
+- pandas, numpy
+
+## Usage
+
+```bash
+pip install RDKit -i https://pypi.tuna.tsinghua.edu.cn/simple
+```
+
+### Basic Run (Domain Generalization Mode, Default)
+
+```bash
+python train_drug_combination_performer_module.py \
+    --checkpoint_path /path/to/performer_module_checkpoint.pt \
+    --hyper_params_path /path/to/hyper_params.pkl \
+    --module_data_dir ../module_data_all \
+    --data_dir ./data_summary
+```
+
+### Custom Paths
+
+```bash
+python train_drug_combination_performer_module.py \
+    --checkpoint_path ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --hyper_params_path ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../module_data_all \
+    --gene_list_path ../../selected_genes_27k.txt\
+    --gene_mapping_path ../gene_mapping.pkl \
+    --gene2vec_path ../../selected_gene2vec_27k.npy \
+    --symbol2ens_path ../../gene2ens.csv \
+    --data_dir ./data_summary \
+    --output_dir ./results \
+    --epochs 500 \
+    --batch_size 256 \
+    --lr 1e-4 \
+    --n_folds 1 \
+    --pool_method max
+```
+
+Reference note: `1e-4 500 0.6824301861702127`.
+
+### Five-Fold Cross-Validation (Disable Domain Generalization)
+
+```bash
+python train_drug_combination_performer_module.py \
+    --checkpoint_path /path/to/checkpoint.pt \
+    --hyper_params_path /path/to/hyper_params.pkl \
+    --module_data_dir ../module_data_all \
+    --no_domain_gen
+```
+
+### Force Reprocessing After Changing the Model or Data
+
+```bash
+python train_drug_combination_performer_module.py \
+    --checkpoint_path /path/to/checkpoint.pt \
+    --hyper_params_path /path/to/hyper_params.pkl \
+    --module_data_dir ../module_data_all \
+    --force_preprocess
+```
+
+## Main Arguments
+
+| Argument | Default | Description |
+| --- | --- | --- |
+| `--checkpoint_path` | Required | Path to the pretrained PerformerModule checkpoint |
+| `--hyper_params_path` | Required | Path to the PerformerModule hyper-parameter pkl file |
+| `--module_data_dir` | None | Directory containing module data files |
+| `--data_dir` | `./data_summary` | Data directory |
+| `--batch_size` | 256 | Training batch size |
+| `--lr` | 1e-4 | Learning rate |
+| `--epochs` | 100 | Number of training epochs |
+| `--n_folds` | 5 | Number of folds |
+| `--domain_gen` | True | Domain generalization mode: train on the full training set and evaluate on the independent test set |
+| `--eval_interval` | 5 | Evaluate every N epochs |
+| `--pool_method` | mean | PerformerModule pooling method (`mean`, `max`, or `concat`) |
+| `--cache_emb` | True | Cache cell embeddings to a `.npy` file |
+| `--force_preprocess` | False | Force regeneration of the PyG dataset |
+
+## Outputs
+
+Results are saved under `--output_dir` (default: `./results/`):
+
+- `best_model_fold{i}.pt` - best model weights for each fold
+- `fpr_tpr_fold{i}.pkl` - FPR/TPR curve data for drawing ROC curves
+- `cell_emb_cache_{pool_method}.npy` - cached PerformerModule cell embeddings
+
+## Evaluation Metrics
+
+Every `--eval_interval` epochs, the script evaluates AUC, PR_AUC, ACC, BACC,
+PREC, TPR, and KAPPA, matching scLong.

BioMod/drug_combination/train_drug_combination_performer_module.py

@@ -0,0 +1,895 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Drug Combination Synergy Prediction with PerformerModule cell embeddings.
+
+Aligned with scLong/drug_combination (DeepDDs) architecture:
+  - Drug branch: PyG GCNConv + global_max_pool (shared weights for drug1 & drug2)
+  - Cell branch: PerformerModule pretrained encoder -> MLP
+  - Fusion: concat(drug1, drug2, cell) -> MLP -> 2-class classification
+
+Usage:
+    python train_drug_combination_performer_module.py \
+        --checkpoint_path <ckpt> --hyper_params_path <pkl> \
+        --module_data_dir ../module_data_all \
+        --data_dir ./data_summary
+"""
+
+import argparse
+import os
+import pickle
+import random
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import networkx as nx
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rdkit import Chem
+from sklearn import metrics
+from tqdm import tqdm
+from sklearn.metrics import (
+    accuracy_score,
+    balanced_accuracy_score,
+    cohen_kappa_score,
+    confusion_matrix,
+    precision_score,
+    recall_score,
+    roc_auc_score,
+)
+from torch.optim.lr_scheduler import StepLR
+
+from torch_geometric import data as DATA
+from torch_geometric.data import InMemoryDataset
+from torch_geometric.loader import DataLoader
+from torch_geometric.nn import GCNConv, global_max_pool as gmp
+
+# ===================== Path setup =====================
+SCRIPT_DIR = Path(__file__).resolve().parent
+PERFORMER_MODULE_ROOT = SCRIPT_DIR.parent
+SCLONG_ROOT = PERFORMER_MODULE_ROOT.parent
+
+if str(PERFORMER_MODULE_ROOT) not in sys.path:
+    sys.path.insert(0, str(PERFORMER_MODULE_ROOT))
+if str(SCLONG_ROOT) not in sys.path:
+    sys.path.insert(0, str(SCLONG_ROOT))
+
+import importlib.util
+
+# Import PerformerModule model
+_models_path = os.path.join(PERFORMER_MODULE_ROOT, 'models_performer_module.py')
+_spec = importlib.util.spec_from_file_location("models_performer_module", _models_path)
+_models_mod = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_models_mod)
+PerformerEncoderOnly_Module = _models_mod.PerformerEncoderOnly_Module
+sanitize_performer_module_config = _models_mod.sanitize_performer_module_config
+
+# Import module_dataset
+_dataset_path = os.path.join(PERFORMER_MODULE_ROOT, 'module_dataset.py')
+_spec2 = importlib.util.spec_from_file_location("module_dataset", _dataset_path)
+_module_mod = importlib.util.module_from_spec(_spec2)
+_spec2.loader.exec_module(_module_mod)
+load_module_data = _module_mod.load_module_data
+
+
+# ===================== Utilities =====================
+
+def print_sys(*args, **kwargs):
+    print(*args, **kwargs, flush=True)
+
+
+def reindex_tensor_universal(tensor, index_positions, dim, filler="0", device="cpu"):
+    index_tensor = torch.tensor(index_positions, device=device)
+    tensor_shape = list(tensor.shape)
+    tensor_shape[dim] = len(index_positions)
+    index_shape = [1] * len(tensor_shape)
+    index_shape[dim] = len(index_positions)
+    index_tensor = index_tensor.view(*index_shape).expand(*tensor_shape)
+    if filler == "0":
+        padder_shape = list(tensor_shape)
+        padder_shape[dim] = 1
+        padder = torch.zeros(padder_shape, dtype=tensor.dtype, device=device)
+    elif filler == "mean":
+        padder = tensor.mean(dim=dim, keepdim=True)
+    else:
+        raise ValueError("filler should be '0' or 'mean'")
+    expanded_tensor = torch.cat((tensor, padder), dim=dim)
+    reindexed = torch.gather(expanded_tensor, dim, index_tensor)
+    return reindexed
+
+
+# ===================== PerformerModule Loading =====================
+
+def load_performer_module_model(
+    checkpoint_path, hyper_params_path, module_data_dir, module_data_path,
+    gene_list_path, gene_mapping_path, gene2vec_path, device,
+    shuffle_module_genes=False, shuffle_seed=2021,
+):
+    print_sys(f"Loading PerformerModule model from: {checkpoint_path}")
+
+    with open(hyper_params_path, 'rb') as f:
+        model_params = pickle.load(f)
+    model_params.pop('module_embed_chunk_padded_gene_budget', None)
+    model_params = sanitize_performer_module_config(
+        model_params, source=hyper_params_path, log_fn=print_sys,
+    )
+
+    if gene2vec_path:
+        model_params['gene2vec_file'] = gene2vec_path
+    model_params['use_gene2vec'] = model_params.get('use_gene2vec', True)
+
+    if 'module_sets' not in model_params or model_params.get('module_sets') is None:
+        if module_data_dir or module_data_path:
+            module_sets = load_module_data(
+                module_data_path=module_data_path,
+                module_data_dir=module_data_dir,
+                gene_list_path=gene_list_path,
+                gene_mapping_path=gene_mapping_path,
+                verbose=True,
+            )
+            model_params['module_sets'] = module_sets
+        else:
+            raise ValueError(
+                "Model needs module_sets but neither --module_data_dir "
+                "nor --module_data_path provided"
+            )
+
+    if shuffle_module_genes:
+        print_sys(f"Shuffling genes across modules with seed={shuffle_seed}")
+        all_gene_indices = []
+        module_sizes = []
+        for module in model_params['module_sets']:
+            all_gene_indices.extend(module)
+            module_sizes.append(len(module))
+        rng = random.Random(shuffle_seed)
+        rng.shuffle(all_gene_indices)
+        shuffled_module_sets = []
+        start_idx = 0
+        for size in module_sizes:
+            shuffled_module_sets.append(all_gene_indices[start_idx:start_idx + size])
+            start_idx += size
+        model_params['module_sets'] = shuffled_module_sets
+        print_sys(f"Shuffled {len(all_gene_indices)} genes across {len(shuffled_module_sets)} modules")
+
+    model = PerformerEncoderOnly_Module(**model_params).to(device)
+
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    state_dict = checkpoint.get('model_state_dict', checkpoint)
+    filtered = {
+        k: v for k, v in state_dict.items()
+        if 'go_conv' not in k and 'G_go' not in k.lower()
+    }
+    model.load_state_dict(filtered, strict=True)
+    print_sys("Loaded pretrained PerformerModule weights.")
+
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad = False
+
+    return model, len(model_params['module_sets']), model_params.get('base_dim', 1024)
+
+
+def build_performer_module_gene_mapping(symbol2ens_path, gene_list_path, expr_genes):
+    gene2ens = pd.read_csv(symbol2ens_path, index_col=0)
+    symbol2ens = gene2ens["Ensembl gene ID"].to_dict()
+
+    input_genes = [symbol2ens.get(sym, sym) for sym in expr_genes]
+    input_gene_num = len(input_genes)
+    input_mapping = {gid: i for i, gid in enumerate(input_genes)}
+
+    with open(gene_list_path, "r") as f:
+        performer_genes = [line.rstrip("\n") for line in f.readlines()]
+    performer_genes_pad = performer_genes + ["PAD"]
+
+    performer_index_positions = [
+        input_mapping.get(gid, input_gene_num) for gid in performer_genes_pad
+    ]
+    intersect = np.intersect1d(input_genes, performer_genes)
+    print_sys(f"Target and PerformerModule genes intersect: {len(intersect)}")
+    return input_mapping, performer_index_positions
+
+
+def build_cell_embeddings(
+    cell_features_df, symbol2ens_path, gene_list_path, gene_mapping_path,
+    gene2vec_path, checkpoint_path, hyper_params_path, module_data_dir,
+    module_data_path, device, pool_method="mean",
+    shuffle_module_genes=False, shuffle_seed=2021,
+):
+    model, num_modules, encoder_embed_dim = load_performer_module_model(
+        checkpoint_path, hyper_params_path, module_data_dir, module_data_path,
+        gene_list_path, gene_mapping_path, gene2vec_path, device,
+        shuffle_module_genes=shuffle_module_genes, shuffle_seed=shuffle_seed,
+    )
+
+    gene_columns = list(cell_features_df.columns)
+    _, performer_index_positions = build_performer_module_gene_mapping(
+        symbol2ens_path, gene_list_path, gene_columns,
+    )
+
+    cell_emb_dict: Dict[str, np.ndarray] = {}
+    first_emb_dim = -1
+
+    with torch.no_grad():
+        for cell_line_id, row in cell_features_df.iterrows():
+            x = torch.tensor(row.values[None, :], dtype=torch.float32, device=device)
+            x_performer = reindex_tensor_universal(
+                x, performer_index_positions, dim=1, filler="0", device=device,
+            )
+            out = model(x_performer, mask_ratio=0.0, return_encodings=True)
+
+            if 'module_encodings' in out:
+                module_enc = out['module_encodings']
+                if pool_method == "mean":
+                    cell_emb = module_enc.mean(dim=1)
+                elif pool_method == "max":
+                    cell_emb = module_enc.max(dim=1)[0]
+                elif pool_method == "concat":
+                    cell_emb = module_enc.reshape(1, -1)
+                else:
+                    cell_emb = module_enc.mean(dim=1)
+            else:
+                raise ValueError(
+                    f"Expected 'module_encodings' in output, got: {list(out.keys())}"
+                )
+
+            cell_emb_np = cell_emb.squeeze(0).cpu().numpy()
+            cell_emb_dict[cell_line_id] = cell_emb_np
+            if first_emb_dim < 0:
+                first_emb_dim = cell_emb_np.shape[0]
+
+    print_sys(
+        f"Built PerformerModule embeddings for {len(cell_emb_dict)} cell lines, "
+        f"dim={first_emb_dim}, pool_method={pool_method}"
+    )
+    return cell_emb_dict, first_emb_dim
+
+
+# ===================== Drug Graph (scLong-aligned: edge_index format) =====================
+
+def one_of_k_encoding(x, allowable_set):
+    if x not in allowable_set:
+        raise Exception(f"input {x} not in allowable set {allowable_set}")
+    return list(map(lambda s: x == s, allowable_set))
+
+
+def one_of_k_encoding_unk(x, allowable_set):
+    if x not in allowable_set:
+        x = allowable_set[-1]
+    return list(map(lambda s: x == s, allowable_set))
+
+
+def atom_features(atom):
+    return np.array(
+        one_of_k_encoding_unk(atom.GetSymbol(), [
+            'C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca',
+            'Fe', 'As', 'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag',
+            'Pd', 'Co', 'Se', 'Ti', 'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd',
+            'In', 'Mn', 'Zr', 'Cr', 'Pt', 'Hg', 'Pb', 'Unknown',
+        ]) +
+        one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +
+        one_of_k_encoding_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +
+        one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +
+        [atom.GetIsAromatic()]
+    )
+
+
+ATOM_FEAT_DIM = 78  # 44 + 11 + 11 + 11 + 1
+
+
+def smile_to_graph(smile):
+    """Convert SMILES to (c_size, features, edge_index) — identical to scLong's creat_data_DC.py."""
+    mol = Chem.MolFromSmiles(smile)
+    if mol is None:
+        return None
+
+    c_size = mol.GetNumAtoms()
+
+    features = []
+    for atom in mol.GetAtoms():
+        feature = atom_features(atom)
+        features.append(feature / (sum(feature) + 1e-8))
+
+    edges = []
+    for bond in mol.GetBonds():
+        edges.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])
+    g = nx.Graph(edges).to_directed()
+    edge_index = []
+    for e1, e2 in g.edges:
+        edge_index.append([e1, e2])
+
+    return c_size, features, edge_index
+
+
+# ===================== TestbedDataset (PyG InMemoryDataset, scLong-aligned) =====================
+
+class TestbedDataset(InMemoryDataset):
+    """
+    PyG InMemoryDataset for drug combination prediction.
+    Same pattern as scLong/drug_combination/utils_test.py::TestbedDataset,
+    but uses PerformerModule cell embeddings instead of scfm/geneformer.
+    """
+
+    def __init__(
+        self, root='data_summary', dataset='drug1',
+        xd=None, xt=None, y=None,
+        cell_emb_dict=None, smile_graph=None,
+        transform=None, pre_transform=None,
+    ):
+        self.dataset = dataset
+        self._xd = xd
+        self._xt = xt
+        self._y = y
+        self._cell_emb_dict = cell_emb_dict
+        self._smile_graph = smile_graph
+        super(TestbedDataset, self).__init__(root, transform, pre_transform)
+
+        if os.path.isfile(self.processed_paths[0]):
+            print_sys(f'Loading preprocessed data: {self.processed_paths[0]}')
+            self.data, self.slices = torch.load(self.processed_paths[0], weights_only=False)
+        else:
+            print_sys(f'Preprocessing data: {self.processed_paths[0]}')
+            self.process(xd, xt, y, cell_emb_dict, smile_graph)
+            self.data, self.slices = torch.load(self.processed_paths[0], weights_only=False)
+
+    @property
+    def raw_file_names(self):
+        return []
+
+    @property
+    def processed_file_names(self):
+        return [self.dataset + '.pt']
+
+    def download(self):
+        pass
+
+    def _download(self):
+        pass
+
+    def _process(self):
+        if not os.path.exists(self.processed_dir):
+            os.makedirs(self.processed_dir)
+
+    def process(self, xd, xt, y, cell_emb_dict, smile_graph):
+        assert len(xd) == len(xt) == len(y), "xd, xt, y must have the same length"
+        data_list = []
+        skipped = 0
+
+        for i in range(len(xd)):
+            smiles = xd[i]
+            cell_name = xt[i]
+            label = y[i]
+
+            if smiles not in smile_graph:
+                skipped += 1
+                continue
+            if cell_name not in cell_emb_dict:
+                skipped += 1
+                continue
+
+            c_size, features, edge_index = smile_graph[smiles]
+
+            GCNData = DATA.Data(
+                x=torch.Tensor(features),
+                edge_index=torch.LongTensor(edge_index).transpose(1, 0),
+                y=torch.Tensor([label]),
+            )
+            cell_emb = cell_emb_dict[cell_name]
+            GCNData.cell = torch.FloatTensor([cell_emb])
+            GCNData.__setitem__('c_size', torch.LongTensor([c_size]))
+            data_list.append(GCNData)
+
+        if skipped > 0:
+            print_sys(f'Skipped {skipped} samples (missing graph or cell embedding)')
+
+        if self.pre_filter is not None:
+            data_list = [d for d in data_list if self.pre_filter(d)]
+        if self.pre_transform is not None:
+            data_list = [self.pre_transform(d) for d in data_list]
+
+        print_sys(f'Processed {len(data_list)} samples. Saving to {self.processed_paths[0]}')
+        data, slices = self.collate(data_list)
+        torch.save((data, slices), self.processed_paths[0])
+
+
+# ===================== Model: GCNNet (PyG, scLong-aligned) =====================
+
+class GCNNet(torch.nn.Module):
+    """
+    GCN-based drug combination model aligned with scLong/drug_combination/models/gcn.py.
+
+    Drug branch: 3-layer GCNConv (78->78->156->312) + global_max_pool + 2 FC layers
+    Drug1 and drug2 share weights (drug2 uses drug1_conv* layers, same as scLong).
+    Cell branch: PerformerModule embedding -> MLP -> output_dim
+    Fusion: concat(drug1, drug2, cell) -> fc1(256) -> fc2(128) -> out(2)
+    """
+
+    def __init__(
+        self, n_output=2, num_features_xd=78, cell_emb_dim=1024,
+        output_dim=128, dropout=0.2,
+    ):
+        super(GCNNet, self).__init__()
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(dropout)
+        self.n_output = n_output
+
+        # Drug graph branch (shared weights for drug1 & drug2)
+        self.drug1_conv1 = GCNConv(num_features_xd, num_features_xd)
+        self.drug1_conv2 = GCNConv(num_features_xd, num_features_xd * 2)
+        self.drug1_conv3 = GCNConv(num_features_xd * 2, num_features_xd * 4)
+        self.drug1_fc_g1 = nn.Linear(num_features_xd * 4, num_features_xd * 2)
+        self.drug1_fc_g2 = nn.Linear(num_features_xd * 2, output_dim)
+
+        # Cell branch (PerformerModule embedding -> MLP)
+        self.reduction = nn.Sequential(
+            nn.Linear(cell_emb_dim, 512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Linear(256, output_dim),
+        )
+
+        # Fusion layers
+        self.fc1 = nn.Linear(3 * output_dim, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.out = nn.Linear(128, self.n_output)
+
+    def forward(self, data1, data2):
+        x1, edge_index1, batch1, cell = (
+            data1.x, data1.edge_index, data1.batch, data1.cell
+        )
+        x2, edge_index2, batch2 = data2.x, data2.edge_index, data2.batch
+
+        # Drug1 branch
+        x1 = self.relu(self.drug1_conv1(x1, edge_index1))
+        x1 = self.relu(self.drug1_conv2(x1, edge_index1))
+        x1 = self.relu(self.drug1_conv3(x1, edge_index1))
+        x1 = gmp(x1, batch1)
+        x1 = self.relu(self.drug1_fc_g1(x1))
+        x1 = self.dropout(x1)
+        x1 = self.drug1_fc_g2(x1)
+        x1 = self.dropout(x1)
+
+        # Drug2 branch (shared weights — uses drug1_conv* layers, same as scLong)
+        x2 = self.relu(self.drug1_conv1(x2, edge_index2))
+        x2 = self.relu(self.drug1_conv2(x2, edge_index2))
+        x2 = self.relu(self.drug1_conv3(x2, edge_index2))
+        x2 = gmp(x2, batch2)
+        x2 = self.relu(self.drug1_fc_g1(x2))
+        x2 = self.dropout(x2)
+        x2 = self.drug1_fc_g2(x2)
+        x2 = self.dropout(x2)
+
+        # Cell branch
+        cell_vector = F.normalize(cell, 2, 1)
+        cell_vector = self.reduction(cell_vector)
+
+        # Fusion
+        xc = torch.cat((x1, x2, cell_vector), 1)
+        xc = self.relu(self.fc1(xc))
+        xc = self.dropout(xc)
+        xc = self.relu(self.fc2(xc))
+        xc = self.dropout(xc)
+        out = self.out(xc)
+        return out
+
+
+# ===================== Training & Evaluation =====================
+
+def train_epoch(model, device, drug1_loader, drug2_loader, optimizer, epoch):
+    """Train one epoch. Returns average loss."""
+    model.train()
+    total_loss = 0.0
+    n_batches = 0
+    pbar = tqdm(
+        zip(drug1_loader, drug2_loader),
+        total=len(drug1_loader),
+        desc=f'Epoch {epoch}',
+        leave=True,
+    )
+    for data1, data2 in pbar:
+        data1 = data1.to(device)
+        data2 = data2.to(device)
+        y = data1.y.view(-1, 1).long().to(device).squeeze(1)
+        optimizer.zero_grad()
+        output = model(data1, data2)
+        loss = F.cross_entropy(output, y)
+        loss.backward()
+        optimizer.step()
+        total_loss += loss.item()
+        n_batches += 1
+        pbar.set_postfix(loss=f'{loss.item():.4f}')
+    avg_loss = total_loss / max(n_batches, 1)
+    return avg_loss
+
+
+def predict(model, device, drug1_loader, drug2_loader):
+    model.eval()
+    total_preds = torch.Tensor()
+    total_labels = torch.Tensor()
+    total_prelabels = torch.Tensor()
+    print_sys(f'Make prediction for {len(drug1_loader.dataset)} samples...')
+
+    with torch.no_grad():
+        for data in zip(drug1_loader, drug2_loader):
+            data1 = data[0].to(device)
+            data2 = data[1].to(device)
+            output = model(data1, data2)
+            ys = F.softmax(output, 1).to('cpu').data.numpy()
+            predicted_labels = list(map(lambda x: np.argmax(x), ys))
+            predicted_scores = list(map(lambda x: x[1], ys))
+            total_preds = torch.cat((total_preds, torch.Tensor(predicted_scores)), 0)
+            total_prelabels = torch.cat(
+                (total_prelabels, torch.Tensor(predicted_labels)), 0
+            )
+            total_labels = torch.cat((total_labels, data1.y.view(-1, 1).cpu()), 0)
+
+    return (
+        total_labels.numpy().flatten(),
+        total_preds.numpy().flatten(),
+        total_prelabels.numpy().flatten(),
+    )
+
+
+def evaluate_metrics(T, S, Y):
+    """Compute all metrics matching scLong's evaluation."""
+    try:
+        AUC = roc_auc_score(T, S)
+    except Exception:
+        AUC = 0.0
+    precision_arr, recall_arr, _ = metrics.precision_recall_curve(T, S)
+    PR_AUC = metrics.auc(recall_arr, precision_arr)
+    BACC = balanced_accuracy_score(T, Y)
+    tn, fp, fn, tp = confusion_matrix(T, Y).ravel()
+    TPR = tp / (tp + fn)
+    PREC = precision_score(T, Y)
+    ACC = accuracy_score(T, Y)
+    KAPPA = cohen_kappa_score(T, Y)
+    RECALL = recall_score(T, Y)
+    return {
+        'AUC': AUC, 'PR_AUC': PR_AUC, 'ACC': ACC, 'BACC': BACC,
+        'PREC': PREC, 'TPR': TPR, 'KAPPA': KAPPA, 'RECALL': RECALL,
+    }
+
+
+# ===================== Cell Feature Loading =====================
+
+def load_cell_features(cell_file):
+    """Load cell features CSV. First row is header (gene_id), second is transcript_id (skipped)."""
+    df = pd.read_csv(cell_file)
+    df = df.iloc[1:].copy()  # skip transcript_id row
+    df = df.set_index(df.columns[0])
+    df = df.astype(float)
+    # Clean cell names: "X22RV1_PROSTATE" -> "22RV1"
+    df.index = [
+        idx.lstrip('X').split('_')[0] if '_' in idx else idx
+        for idx in df.index
+    ]
+    return df
+
+
+# ===================== Main =====================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Drug Combination Synergy Prediction with PerformerModule'
+    )
+
+    # PerformerModule model args
+    parser.add_argument('--checkpoint_path', type=str, required=True,
+                        help='Path to PerformerModule pretrained checkpoint')
+    parser.add_argument('--hyper_params_path', type=str, required=True,
+                        help='Path to PerformerModule hyper parameters pkl')
+    parser.add_argument('--module_data_dir', type=str, default=None,
+                        help='Directory containing module data files')
+    parser.add_argument('--module_data_path', type=str, default=None,
+                        help='Path to single module data file')
+    parser.add_argument('--gene_list_path', type=str,
+                        default=os.path.join(PERFORMER_MODULE_ROOT, 'gene_list_27k.txt'))
+    parser.add_argument('--gene_mapping_path', type=str,
+                        default=os.path.join(PERFORMER_MODULE_ROOT, 'gene_mapping.pkl'))
+    parser.add_argument('--gene2vec_path', type=str,
+                        default=os.path.join(PERFORMER_MODULE_ROOT, 'gene2vec_27k.npy'))
+    parser.add_argument('--symbol2ens_path', type=str,
+                        default=os.path.join(SCLONG_ROOT, 'gene2ensemble_ID.csv'))
+    parser.add_argument('--pool_method', type=str, default='mean',
+                        choices=['mean', 'max', 'concat'])
+    parser.add_argument('--shuffle_module_genes', action='store_true',
+                        help='Disrupt the gene allocation within the module (used for comparative experiments to verify the effectiveness of module division)')
+    parser.add_argument('--shuffle_seed', type=int, default=2021,
+                        help='The random seed used when scrambling the module gene (needs to be consistent with the value used during pre-training)')
+
+    # Data args
+    parser.add_argument('--data_dir', type=str,
+                        default=os.path.join(SCRIPT_DIR, 'data_summary'))
+    parser.add_argument('--label_file', type=str, default='new_labels_0_10.csv')
+    parser.add_argument('--cell_file', type=str, default='new_cell_features_954.csv')
+    parser.add_argument('--test_label_file', type=str,
+                        default='test_data/independent_input.csv')
+    parser.add_argument('--test_cell_file', type=str,
+                        default='test_data/independent_cell_features_954.csv')
+
+    # Training args
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--n_folds', type=int, default=5)
+    parser.add_argument('--domain_gen', action='store_true', default=True,
+                        help='Domain generalization: train full set, test on independent set')
+    parser.add_argument('--no_domain_gen', dest='domain_gen', action='store_false')
+    parser.add_argument('--eval_interval', type=int, default=5)
+    parser.add_argument('--output_dir', type=str,
+                        default=os.path.join(SCRIPT_DIR, 'results'))
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--cache_emb', action='store_true', default=True,
+                        help='Cache PerformerModule cell embeddings to .npy')
+    parser.add_argument('--no_cache_emb', dest='cache_emb', action='store_false')
+    parser.add_argument('--force_preprocess', action='store_true', default=False,
+                        help='Force re-creation of PyG .pt datasets')
+
+    args = parser.parse_args()
+
+    # Device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print_sys(f'Using device: {device}')
+
+    # Seed
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # -------- Step 1: Build PerformerModule cell embeddings --------
+    train_cell_file = os.path.join(args.data_dir, args.cell_file)
+    train_cell_df = load_cell_features(train_cell_file)
+    print_sys(f'Loaded train cell features: {train_cell_df.shape}')
+
+    test_cell_file = os.path.join(args.data_dir, args.test_cell_file)
+    test_cell_df = load_cell_features(test_cell_file)
+    print_sys(f'Loaded test cell features: {test_cell_df.shape}')
+
+    # Combine all cell lines and remove duplicates
+    all_cell_df = pd.concat([train_cell_df, test_cell_df])
+    all_cell_df = all_cell_df[~all_cell_df.index.duplicated(keep='first')]
+    print_sys(f'Total unique cell lines: {len(all_cell_df)}')
+
+    # Cache embeddings
+    cache_path = os.path.join(args.output_dir, f'cell_emb_cache_{args.pool_method}.npy')
+    cache_dim_path = os.path.join(
+        args.output_dir, f'cell_emb_dim_{args.pool_method}.txt'
+    )
+
+    if args.cache_emb and os.path.exists(cache_path) and not args.force_preprocess:
+        print_sys(f'Loading cached cell embeddings from {cache_path}')
+        cell_emb_dict = np.load(cache_path, allow_pickle=True).item()
+        with open(cache_dim_path) as f:
+            cell_emb_dim = int(f.read().strip())
+    else:
+        cell_emb_dict, cell_emb_dim = build_cell_embeddings(
+            all_cell_df, args.symbol2ens_path, args.gene_list_path,
+            args.gene_mapping_path, args.gene2vec_path, args.checkpoint_path,
+            args.hyper_params_path, args.module_data_dir, args.module_data_path,
+            device, args.pool_method,
+            shuffle_module_genes=args.shuffle_module_genes,
+            shuffle_seed=args.shuffle_seed,
+        )
+        if args.cache_emb:
+            np.save(cache_path, cell_emb_dict)
+            with open(cache_dim_path, 'w') as f:
+                f.write(str(cell_emb_dim))
+            print_sys(f'Cached cell embeddings to {cache_path}')
+
+    # -------- Step 2: Build SMILES graphs --------
+    train_label_path = os.path.join(args.data_dir, args.label_file)
+    train_labels = pd.read_csv(train_label_path)
+    print_sys(f'Train labels: {len(train_labels)} samples')
+
+    test_label_path = os.path.join(args.data_dir, args.test_label_file)
+    test_labels = pd.read_csv(test_label_path)
+    print_sys(f'Test labels: {len(test_labels)} samples')
+
+    all_smiles = set(
+        train_labels['drug1'].tolist() + train_labels['drug2'].tolist()
+        + test_labels['drug1'].tolist() + test_labels['drug2'].tolist()
+    )
+    smile_graph = {}
+    failed = []
+    for smi in all_smiles:
+        result = smile_to_graph(smi)
+        if result is not None:
+            smile_graph[smi] = result
+        else:
+            failed.append(smi)
+    if failed:
+        print_sys(f'Warning: {len(failed)} SMILES failed to parse')
+    print_sys(f'Built graphs for {len(smile_graph)} unique SMILES')
+
+    # -------- Step 3: Create PyG datasets --------
+    # Delete cached .pt files if force_preprocess
+    if args.force_preprocess:
+        processed_dir = os.path.join(args.data_dir, 'processed')
+        if os.path.exists(processed_dir):
+            import shutil
+            shutil.rmtree(processed_dir)
+            print_sys(f'Removed cached processed data: {processed_dir}')
+
+    # Train datasets (drug1 and drug2 separate, scLong pattern)
+    drug1_train = train_labels['drug1'].tolist()
+    drug2_train = train_labels['drug2'].tolist()
+    cells_train = train_labels['cell'].tolist()
+    labels_train = train_labels['label'].tolist()
+
+    train_d1_dataset = TestbedDataset(
+        root=args.data_dir, dataset='train_drug1_performer_module',
+        xd=drug1_train, xt=cells_train, y=labels_train,
+        cell_emb_dict=cell_emb_dict, smile_graph=smile_graph,
+    )
+    train_d2_dataset = TestbedDataset(
+        root=args.data_dir, dataset='train_drug2_performer_module',
+        xd=drug2_train, xt=cells_train, y=labels_train,
+        cell_emb_dict=cell_emb_dict, smile_graph=smile_graph,
+    )
+
+    # Test datasets (independent set)
+    drug1_test = test_labels['drug1'].tolist()
+    drug2_test = test_labels['drug2'].tolist()
+    cells_test = test_labels['cell'].tolist()
+    labels_test = test_labels['label'].tolist()
+
+    test_d1_dataset = TestbedDataset(
+        root=args.data_dir, dataset='test_drug1_performer_module',
+        xd=drug1_test, xt=cells_test, y=labels_test,
+        cell_emb_dict=cell_emb_dict, smile_graph=smile_graph,
+    )
+    test_d2_dataset = TestbedDataset(
+        root=args.data_dir, dataset='test_drug2_performer_module',
+        xd=drug2_test, xt=cells_test, y=labels_test,
+        cell_emb_dict=cell_emb_dict, smile_graph=smile_graph,
+    )
+
+    # -------- Step 4: Training (5-fold, scLong-aligned) --------
+    lenth = len(train_d1_dataset)
+    pot = lenth // args.n_folds
+    random_num = random.sample(range(lenth), lenth)
+    all_fold_aucs = []
+
+    for fold_i in range(args.n_folds):
+        print_sys(f'\n{"=" * 60}')
+        print_sys(f'Fold {fold_i + 1}/{args.n_folds}')
+        print_sys(f'{"=" * 60}')
+
+        if args.domain_gen:
+            # Domain generalization: train on full train set, test on independent set
+            d1_train_data = train_d1_dataset
+            d1_test_data = test_d1_dataset
+            d2_train_data = train_d2_dataset
+            d2_test_data = test_d2_dataset
+        else:
+            # Standard k-fold cross-validation
+            test_num = random_num[pot * fold_i : pot * (fold_i + 1)]
+            train_num = (
+                random_num[: pot * fold_i] + random_num[pot * (fold_i + 1) :]
+            )
+            d1_train_data = train_d1_dataset[train_num]
+            d1_test_data = train_d1_dataset[test_num]
+            d2_train_data = train_d2_dataset[train_num]
+            d2_test_data = train_d2_dataset[test_num]
+
+        d1_loader_train = DataLoader(
+            d1_train_data, batch_size=args.batch_size, shuffle=None, pin_memory=True,
+        )
+        d1_loader_test = DataLoader(
+            d1_test_data, batch_size=args.batch_size, shuffle=None, pin_memory=True,
+        )
+        d2_loader_train = DataLoader(
+            d2_train_data, batch_size=args.batch_size, shuffle=None, pin_memory=True,
+        )
+        d2_loader_test = DataLoader(
+            d2_test_data, batch_size=args.batch_size, shuffle=None, pin_memory=True,
+        )
+
+        model = GCNNet(cell_emb_dim=cell_emb_dim).to(device)
+
+        if fold_i == 0:
+            for name, param in model.named_parameters():
+                print_sys(f"  Parameter: {name}, Shape: {param.size()}")
+            total_params = sum(p.numel() for p in model.parameters())
+            print_sys(f'Total model parameters: {total_params}')
+
+        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+        scheduler = StepLR(optimizer, step_size=1000, gamma=0.9)
+
+        best_auc = 0
+        loss_history = []
+        metric_history = []
+
+        for epoch in range(1, args.epochs + 1):
+            avg_loss = train_epoch(
+                model, device, d1_loader_train, d2_loader_train,
+                optimizer, epoch,
+            )
+            scheduler.step()
+            loss_history.append(avg_loss)
+
+            if epoch % args.eval_interval == 0:
+                T, S, Y = predict(model, device, d1_loader_test, d2_loader_test)
+                m = evaluate_metrics(T, S, Y)
+                metric_history.append({'epoch': epoch, **m})
+
+                if m['AUC'] > best_auc:
+                    best_auc = m['AUC']
+                    best_fpr, best_tpr = metrics.roc_curve(T, S)[:2]
+                    save_path = os.path.join(
+                        args.output_dir, f'best_model_fold{fold_i}.pt'
+                    )
+                    torch.save(model.state_dict(), save_path)
+                    print_sys(f'** New best AUROC: {best_auc:.4f} (saved)')
+
+                print_sys(
+                    f'Epoch {epoch} | '
+                    f'AUROC: {m["AUC"]:.4f}  PR_AUC: {m["PR_AUC"]:.4f}  '
+                    f'ACC: {m["ACC"]:.4f}  BACC: {m["BACC"]:.4f}  '
+                    f'PREC: {m["PREC"]:.4f}  TPR: {m["TPR"]:.4f}  '
+                    f'KAPPA: {m["KAPPA"]:.4f}'
+                )
+
+        # Save loss curve
+        loss_csv_path = os.path.join(args.output_dir, f'loss_fold{fold_i}.csv')
+        loss_df = pd.DataFrame({'epoch': range(1, len(loss_history) + 1), 'avg_loss': loss_history})
+        loss_df.to_csv(loss_csv_path, index=False)
+
+        # Plot loss curve
+        fig, ax = plt.subplots(figsize=(8, 5))
+        ax.plot(loss_df['epoch'], loss_df['avg_loss'], linewidth=1.5)
+        ax.set_xlabel('Epoch')
+        ax.set_ylabel('Average Loss')
+        ax.set_title(f'Training Loss - Fold {fold_i + 1}')
+        ax.grid(True, alpha=0.3)
+        fig.tight_layout()
+        fig.savefig(os.path.join(args.output_dir, f'loss_curve_fold{fold_i}.png'), dpi=150)
+        plt.close(fig)
+
+        # Save metrics history
+        if metric_history:
+            metric_df = pd.DataFrame(metric_history)
+            metric_df.to_csv(os.path.join(args.output_dir, f'metrics_fold{fold_i}.csv'), index=False)
+
+        # Plot ROC curve (using best model's predictions)
+        fig, ax = plt.subplots(figsize=(6, 6))
+        ax.plot(best_fpr, best_tpr, linewidth=2, label=f'AUROC = {best_auc:.4f}')
+        ax.plot([0, 1], [0, 1], 'k--', linewidth=1, alpha=0.5)
+        ax.set_xlim([0, 1])
+        ax.set_ylim([0, 1.02])
+        ax.set_xlabel('False Positive Rate')
+        ax.set_ylabel('True Positive Rate')
+        ax.set_title(f'ROC Curve - Fold {fold_i + 1}')
+        ax.legend(loc='lower right')
+        ax.grid(True, alpha=0.3)
+        fig.tight_layout()
+        fig.savefig(os.path.join(args.output_dir, f'roc_curve_fold{fold_i}.png'), dpi=150)
+        plt.close(fig)
+
+        # Save fpr/tpr data
+        fpr_tpr_path = os.path.join(args.output_dir, f'fpr_tpr_fold{fold_i}.pkl')
+        with open(fpr_tpr_path, 'wb') as fp:
+            pickle.dump((best_fpr, best_tpr), fp)
+
+        print_sys(f'Fold {fold_i + 1} best AUROC: {best_auc:.4f}')
+        all_fold_aucs.append(best_auc)
+
+    # Summary
+    print_sys(f'\n{"=" * 60}')
+    print_sys(f'All folds AUROC: {all_fold_aucs}')
+    print_sys(f'Mean AUROC: {np.mean(all_fold_aucs):.4f} +/- {np.std(all_fold_aucs):.4f}')
+    print_sys('Training complete.')
+
+
+if __name__ == '__main__':
+    main()

BioMod/drug_response/cmd.txt

@@ -0,0 +1,10 @@
+python train_drug_performer_module.py \
+    --gpu_id 0 \
+    --data_path ../../drug_response_pytorch/data \
+    --checkpoint_path ../performer_module_multiply_all_module_48Mcell/performer_module_multiply_all_module_48Mcell_epoch2.pth \
+    --hyper_params_path ../performer_module_multiply_all_module_48Mcell/performer_module_multiply_all_module_48Mcell.pkl \
+    --epochs 300 \
+    --module_data_dir ../module_data_all \
+    --save_dir ./performer_module_multiply_all_module_48Mcell \
+    --fusion_method multiply \
+    --pool_method max

BioMod/drug_response/train_drug_performer_module.py

@@ -0,0 +1,1065 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Drug response prediction: using drug molecular structure + PerformerModule cell embedding
+
+Architecture:
+- Drug branch: 3-layer GCN + residual connection + double pooling
+- Cell branch: extract module-level features from pre-trained PerformerModule
+- Fusion layer: attention mechanism + deep MLP
+- Optimization strategy: AdamW + cosine annealing + early stopping
+
+How to use:
+    python train_drug_performer_module.py --gpu_id 0 --data_path ../drug_response_pytorch/data --epochs 300
+"""
+
+import argparse
+import os
+import random
+import sys
+import pickle
+from typing import Dict, List, Tuple, Union, Optional
+
+import hickle as hkl
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+import torch.nn as nn
+from scipy.stats import pearsonr
+from torch.utils.data import DataLoader, Dataset
+
+# Import PerformerModule model
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+sys.path.insert(0, os.path.dirname(project_root))  # Add scLong root directory
+
+# Use direct path import to avoid conflicts with files with the same name
+import importlib.util
+performer_module_dir = project_root
+
+# import models_performer_module
+models_path = os.path.join(performer_module_dir, 'models_performer_module.py')
+if not os.path.exists(models_path):
+    raise FileNotFoundError(f"Model file not found: {models_path}")
+spec = importlib.util.spec_from_file_location("models_performer_module", models_path)
+models_performer_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(models_performer_module)
+PerformerEncoderOnly_Module = models_performer_module.PerformerEncoderOnly_Module
+sanitize_performer_module_config = models_performer_module.sanitize_performer_module_config
+
+# Import module_dataset
+dataset_path = os.path.join(performer_module_dir, 'module_dataset.py')
+if not os.path.exists(dataset_path):
+    raise FileNotFoundError(f"Dataset file not found: {dataset_path}")
+spec = importlib.util.spec_from_file_location("module_dataset", dataset_path)
+module_dataset_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(module_dataset_module)
+load_module_data = module_dataset_module.load_module_data
+
+
+# ===================== Utility functions =====================
+
+def print_sys(*args, **kwargs):
+    """Print and refresh output"""
+    print(*args, **kwargs, flush=True)
+
+
+def reindex_tensor_universal(
+    tensor: torch.Tensor,
+    index_positions: List[int],
+    dim: int,
+    filler: str = "0",
+    device: Union[torch.device, str] = "cpu",
+) -> torch.Tensor:
+    """Universal reordering operation"""
+    index_tensor = torch.tensor(index_positions, device=device)
+    tensor_shape = list(tensor.shape)
+    tensor_shape[dim] = len(index_positions)
+    index_shape = [1] * len(tensor_shape)
+    index_shape[dim] = len(index_positions)
+    index_tensor = index_tensor.view(*index_shape).expand(*tensor_shape)
+
+    if filler == "0":
+        padder_shape = list(tensor_shape)
+        padder_shape[dim] = 1
+        padder = torch.zeros(padder_shape, dtype=tensor.dtype, device=device)
+    elif filler == "mean":
+        padder = tensor.mean(dim=dim, keepdim=True)
+    else:
+        raise ValueError("filler should be '0' or 'mean'")
+
+    expanded_tensor = torch.cat((tensor, padder), dim=dim)
+    reindexed = torch.gather(expanded_tensor, dim, index_tensor)
+    return reindexed
+
+
+# ===================== PerformerModule related =====================
+
+def load_performer_module_model(
+    checkpoint_path: str,
+    hyper_params_path: str,
+    module_data_dir: Optional[str],
+    module_data_path: Optional[str],
+    gene_list_path: str,
+    gene_mapping_path: str,
+    gene2vec_path: str,
+    device: torch.device,
+    shuffle_module_genes: bool = False,
+    seed: int = 0,
+) -> PerformerEncoderOnly_Module:
+    """Load a pretrained PerformerModule model"""
+    print_sys(f"Loading PerformerModule model from: {checkpoint_path}")
+
+    # Load model hyperparameters
+    print_sys(f"Loading hyper parameters from: {hyper_params_path}")
+    with open(hyper_params_path, 'rb') as f:
+        model_params = pickle.load(f)
+    model_params.pop('module_embed_chunk_padded_gene_budget', None)
+    model_params = sanitize_performer_module_config(
+        model_params,
+        source=hyper_params_path,
+        log_fn=print_sys,
+    )
+
+    # Update gene2vec path
+    if gene2vec_path:
+        model_params['gene2vec_file'] = gene2vec_path
+
+    # Read fusion_method from pkl to ensure consistency with pre-training
+    fusion_method = model_params.get('fusion_method', 'multiply')
+    print_sys(f"Using fusion_method from pkl: {fusion_method}")
+
+    # use_gene2vec strictly follows the pre-training pkl to avoid conflicts caused by downstream coverage.
+    pkl_use_gene2vec = model_params.get('use_gene2vec', True)
+    model_params['use_gene2vec'] = pkl_use_gene2vec
+    print_sys(f"use_gene2vec={pkl_use_gene2vec} (from PKL)")
+    
+    # Load module data (if not present in hyperparameters)
+    if 'module_sets' not in model_params or model_params.get('module_sets') is None:
+        if module_data_dir or module_data_path:
+            print_sys("Loading module data...")
+            module_sets = load_module_data(
+                module_data_path=module_data_path,
+                module_data_dir=module_data_dir,
+                gene_list_path=gene_list_path,
+                gene_mapping_path=gene_mapping_path,
+                verbose=True
+            )
+            model_params['module_sets'] = module_sets
+        else:
+            raise ValueError("Model requires module_sets, but --module_data_dir or --module_data_path was not provided")
+    else:
+        module_sets = model_params['module_sets']
+        print_sys(f"Using module_sets from hyper parameters ({len(module_sets)} modules)")
+
+    # If shuffle_module_genes is enabled, disrupt the gene distribution within the module.
+    if shuffle_module_genes:
+        print_sys(
+            f"Applying downstream-only extra module reshuffle (seed={seed}). "
+            f"Pretraining-time module_sets from PKL will be reshuffled again."
+        )
+        # Collect all gene indices
+        all_gene_indices = []
+        module_sizes = []
+        for module in module_sets:
+            all_gene_indices.extend(module)
+            module_sizes.append(len(module))
+
+        # Scramble the gene index using a fixed random seed
+        rng = random.Random(seed)
+        rng.shuffle(all_gene_indices)
+
+        # Redistribute genes according to original module size
+        shuffled_module_sets = []
+        start_idx = 0
+        for size in module_sizes:
+            shuffled_module_sets.append(all_gene_indices[start_idx:start_idx + size])
+            start_idx += size
+        module_sets = shuffled_module_sets
+        model_params['module_sets'] = module_sets
+        print_sys(f"Shuffled {len(all_gene_indices)} genes across {len(module_sets)} modules")
+
+    print_sys(f"PerformerModule config: base_dim={model_params.get('base_dim', 'N/A')}, "
+              f"num_modules={len(module_sets)}")
+    
+    # Build model
+    model = PerformerEncoderOnly_Module(**model_params).to(device)
+    
+    # Load weights
+    print_sys(f"Loading checkpoint from: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    
+    # Handle parameters in checkpoint
+    if 'model_state_dict' in checkpoint:
+        state_dict = checkpoint['model_state_dict']
+    else:
+        state_dict = checkpoint
+    
+    # Filter out possible GO related parameters
+    filtered_state_dict = {}
+    for k, v in state_dict.items():
+        if 'go_conv' not in k and 'G_go' not in k.lower():
+            filtered_state_dict[k] = v
+
+    # Use strict=True to ensure the weights match exactly (except for filtered GO parameters)
+    model.load_state_dict(filtered_state_dict, strict=True)
+    
+    print_sys("Loaded pretrained PerformerModule weights.")
+    
+    # freeze model
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad = False
+    
+    return model, len(module_sets), model_params.get('base_dim', 1024)
+
+
+def build_performer_module_gene_mapping(
+    symbol2ens_path: str,
+    gene_list_path: str,
+    expr_genes: List[str],
+) -> Tuple[Dict[str, int], List[int]]:
+    """Build a gene map"""
+    gene2ens = pd.read_csv(symbol2ens_path, index_col=0)
+    symbol2ens = gene2ens["Ensembl gene ID"].to_dict()
+
+    input_genes = [symbol2ens.get(sym, sym) for sym in expr_genes]
+    input_gene_num = len(input_genes)
+    input_mapping = {gid: i for i, gid in enumerate(input_genes)}
+
+    with open(gene_list_path, "r") as f:
+        performer_genes = [line.rstrip("\n") for line in f.readlines()]
+    performer_genes_pad = performer_genes + ["PAD"]
+
+    performer_index_positions = [
+        input_mapping.get(gid, input_gene_num) for gid in performer_genes_pad
+    ]
+    intersect = np.intersect1d(input_genes, performer_genes)
+    print_sys(f"Target and PerformerModule genes intersect: {len(intersect)}")
+    return input_mapping, performer_index_positions
+
+
+def build_cell_embeddings_with_performer_module(
+    gexpr_feature: pd.DataFrame,
+    symbol2ens_path: str,
+    gene_list_path: str,
+    gene_mapping_path: str,
+    gene2vec_path: str,
+    checkpoint_path: str,
+    hyper_params_path: str,
+    module_data_dir: Optional[str],
+    module_data_path: Optional[str],
+    device: torch.device,
+    pool_method: str = "mean",
+    shuffle_module_genes: bool = False,
+    seed: int = 0,
+) -> Tuple[Dict[str, np.ndarray], int]:
+    """
+    Use PerformerModule to construct embedding for each cell line
+
+    Args:
+        pool_method: pooling method, 'mean', 'max', 'cls', or 'concat'
+            - 'mean': average all modules (default)
+            - 'max': take the maximum value for all modules
+            - 'cls': only use CLS token (if available)
+            - 'concat': concatenate all modules (the dimension will be very large)
+        use_gene2vec: automatically restored by pre-trained PKL, no longer covered downstream
+    """
+    model, num_modules, encoder_embed_dim = load_performer_module_model(
+        checkpoint_path=checkpoint_path,
+        hyper_params_path=hyper_params_path,
+        module_data_dir=module_data_dir,
+        module_data_path=module_data_path,
+        gene_list_path=gene_list_path,
+        gene_mapping_path=gene_mapping_path,
+        gene2vec_path=gene2vec_path,
+        device=device,
+        shuffle_module_genes=shuffle_module_genes,
+        seed=seed,
+    )
+
+    _, performer_index_positions = build_performer_module_gene_mapping(
+        symbol2ens_path=symbol2ens_path,
+        gene_list_path=gene_list_path,
+        expr_genes=list(gexpr_feature.columns),
+    )
+
+    cell_emb_dict: Dict[str, np.ndarray] = {}
+    first_emb_dim: int = -1
+
+    with torch.no_grad():
+        for cell_line_id, row in gexpr_feature.iterrows():
+            x = torch.tensor(row.values[None, :], dtype=torch.float32, device=device)
+            x_performer = reindex_tensor_universal(
+                x, performer_index_positions, dim=1, filler="0", device=device
+            )
+
+            # Extract features through encoder (without masking, mask_ratio=0.0)
+            out = model(x_performer, mask_ratio=0.0, return_encodings=True)
+            
+            # Extract module encodings
+            if 'module_encodings' in out:
+                module_encodings = out['module_encodings']  # (1, num_modules, encoder_embed_dim)
+                
+                # Processed according to pooling method
+                if pool_method == "cls":
+                    if model.use_cls:
+                        # Use CLS token (in latent, but module_encodings does not contain CLS)
+                        # Need to get the complete latent from forward_encoder
+                        enc_out = model.forward_encoder(x_performer, mask_ratio=0.0)
+                        latent = enc_out[0]
+                        if model.use_cls:
+                            cell_emb = latent[:, 0, :]  # (1, encoder_embed_dim)
+                        else:
+                            print_sys("Warning: pool_method='cls' but model has no CLS token, using 'mean' instead")
+                            cell_emb = module_encodings.mean(dim=1)
+                    else:
+                        print_sys("Warning: pool_method='cls' but model has no CLS token, using 'mean' instead")
+                        cell_emb = module_encodings.mean(dim=1)
+                elif pool_method == "mean":
+                    cell_emb = module_encodings.mean(dim=1)  # (1, encoder_embed_dim)
+                elif pool_method == "max":
+                    cell_emb = module_encodings.max(dim=1)[0]  # (1, encoder_embed_dim)
+                elif pool_method == "concat":
+                    cell_emb = module_encodings.reshape(1, -1)  # (1, num_modules * encoder_embed_dim)
+                else:
+                    raise ValueError(f"Unknown pool_method: {pool_method}")
+            else:
+                raise ValueError(f"Expected 'module_encodings' in output, got: {list(out.keys())}")
+            
+            cell_emb_np = cell_emb.squeeze(0).cpu().numpy()
+            cell_emb_dict[cell_line_id] = cell_emb_np
+            if first_emb_dim < 0:
+                first_emb_dim = cell_emb_np.shape[0]
+
+    print_sys(
+        f"Built PerformerModule embeddings for {len(cell_emb_dict)} cell lines, "
+        f"dim={first_emb_dim}, pool_method={pool_method}"
+    )
+    return cell_emb_dict, first_emb_dim
+
+
+# ===================== Data processing =====================
+
+TCGA_LABEL_SET = [
+    "ALL", "BLCA", "BRCA", "CESC", "DLBC", "LIHC", "LUAD", "ESCA", "GBM",
+    "HNSC", "KIRC", "LAML", "LCML", "LGG", "LUSC", "MESO", "MM", "NB",
+    "OV", "PAAD", "SCLC", "SKCM", "STAD", "THCA", "COAD/READ",
+]
+
+
+def normalize_adj(adj: np.ndarray) -> np.ndarray:
+    """Normalized adjacency matrix"""
+    adj = adj + np.eye(adj.shape[0])
+    d = sp.diags(np.power(np.array(adj.sum(1)), -0.5).flatten(), 0).toarray()
+    a_norm = adj.dot(d).transpose().dot(d)
+    return a_norm
+
+
+def calculate_graph_feat(
+    feat_mat: np.ndarray,
+    adj_list: List[List[int]],
+    max_atoms: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute drug graph features"""
+    assert feat_mat.shape[0] == len(adj_list)
+    feat = np.zeros((max_atoms, feat_mat.shape[-1]), dtype=np.float32)
+    adj_mat = np.zeros((max_atoms, max_atoms), dtype=np.float32)
+
+    feat[: feat_mat.shape[0], :] = feat_mat
+    for i, nodes in enumerate(adj_list):
+        for each in nodes:
+            adj_mat[i, int(each)] = 1
+
+    assert np.allclose(adj_mat, adj_mat.T)
+    adj_ = adj_mat[: len(adj_list), : len(adj_list)]
+    norm_adj_ = normalize_adj(adj_)
+    adj_mat[: len(adj_list), : len(adj_list)] = norm_adj_
+    return feat, adj_mat
+
+
+def metadata_generate(
+    drug_info_file: str,
+    cell_line_info_file: str,
+    drug_feature_dir: str,
+    gene_expression_file: str,
+) -> Tuple[dict, pd.DataFrame, list]:
+    """Generate metadata"""
+    # drug_id --> pubchem_id
+    reader = pd.read_csv(drug_info_file)
+    drugid2pubchemid = {# Build drug_id → pubchem_id mapping
+        str(row.iloc[0]): str(row.iloc[5])# Mapping from column 0 of drug information CSV to column 5 (pubchem_id)
+        for _, row in reader.iterrows()
+        if str(row.iloc[5]).isdigit()
+    }
+
+    # cellline --> cancer type # Read from the TSV file, the first column is cellline_id, the last column is the cancer type
+    cellline2cancertype: Dict[str, str] = {}
+    with open(cell_line_info_file, "r") as f:
+        for line in f.readlines()[1:]:
+            cellline_id = line.split("\t")[1]
+            tcga_label = line.strip().split("\t")[-1]
+            cellline2cancertype[cellline_id] = tcga_label
+
+    gexpr_feature = pd.read_csv(# Gene expression matrix for cell lines
+        gene_expression_file, sep=",", header=0, index_col=0
+    )
+
+    # load drug graph features
+    drug_feature: Dict[str, list] = {}
+    drug_pubchem_id_set = []
+    for fname in os.listdir(drug_feature_dir):
+        if not fname.endswith(".hkl"):
+            continue
+        pubchem_id = fname.split(".")[0]
+        feat_mat, adj_list, degree_list = hkl.load(
+            os.path.join(drug_feature_dir, fname)
+        )
+        drug_feature[pubchem_id] = [feat_mat, adj_list, degree_list]# [Feature matrix, adjacency list, degree list]
+        drug_pubchem_id_set.append(pubchem_id)
+
+    # IC50
+     # Read IC50 response data from GDSC database
+    # Filter out drug-cell line pairs with corresponding pubchem_id and gene expression data
+    cancer_response_exp_file = os.path.join(
+        os.path.dirname(gene_expression_file), "GDSC_IC50.csv"
+    )
+    experiment_data = pd.read_csv(
+        cancer_response_exp_file, sep=",", header=0, index_col=0
+    )
+    drug_match_list = [
+        item
+        for item in experiment_data.index
+        if item.split(":")[1] in drugid2pubchemid.keys()
+    ]
+    experiment_data_filtered = experiment_data.loc[drug_match_list]
+
+    data_idx: list = []# Each entry contains: cell line ID, drug pubchem ID, IC50 value, cancer type.
+    for each_drug in experiment_data_filtered.index:
+        for each_cellline in experiment_data_filtered.columns:
+            pubchem_id = drugid2pubchemid[each_drug.split(":")[-1]]
+            if (
+                str(pubchem_id) in drug_pubchem_id_set
+                and each_cellline in gexpr_feature.index
+            ):
+                if (
+                    not np.isnan(
+                        experiment_data_filtered.loc[each_drug, each_cellline]
+                    )
+                    and each_cellline in cellline2cancertype.keys()
+                ):
+                    ln_ic50 = float(
+                        experiment_data_filtered.loc[each_drug, each_cellline]
+                    )
+                    data_idx.append(
+                        (
+                            each_cellline,
+                            pubchem_id,
+                            ln_ic50,
+                            cellline2cancertype[each_cellline],
+                        )
+                    )
+
+    nb_celllines = len(set([item[0] for item in data_idx]))
+    nb_drugs = len(set([item[1] for item in data_idx]))
+    print_sys(
+        f"{len(data_idx)} instances across {nb_celllines} cell lines and {nb_drugs} drugs."
+    )
+    return drug_feature, gexpr_feature, data_idx # Drug molecular graph features, cell line gene expression features, effective (cell line, drug, IC50, cancer type) four-tuple list
+
+
+def data_split(data_idx: list, ratio: float = 0.95):
+    """Stratified by cancer type"""
+    train_idx, test_idx = [], []
+    for each_type in TCGA_LABEL_SET:
+        data_subtype_idx = [item for item in data_idx if item[-1] == each_type]
+        if len(data_subtype_idx) == 0:
+            continue
+        k = max(1, int(ratio * len(data_subtype_idx)))
+        train_list = random.sample(data_subtype_idx, k)
+        test_list = [item for item in data_subtype_idx if item not in train_list]
+        train_idx += train_list
+        test_idx += test_list
+    return train_idx, test_idx
+
+
+def feature_extract(
+    data_idx: list,
+    drug_feature: dict,
+    max_atoms: int,
+) -> Tuple[list, np.ndarray, list, list]:
+    """Extract drug map features"""
+    nb_instance = len(data_idx)
+
+    drug_data = [[] for _ in range(nb_instance)]
+    target = np.zeros(nb_instance, dtype=np.float32)
+    cancer_type_list = []
+
+    for idx in range(nb_instance):
+        cell_line_id, pubchem_id, ln_ic50, cancer_type = data_idx[idx]
+
+        # Drug map features
+        feat_mat, adj_list, _ = drug_feature[str(pubchem_id)]
+        feat, adj = calculate_graph_feat(feat_mat, adj_list, max_atoms)
+        drug_data[idx] = [feat, adj]
+
+        # target value
+        target[idx] = ln_ic50
+
+        cancer_type_list.append([cancer_type, cell_line_id, pubchem_id])
+
+    return drug_data, target, cancer_type_list, data_idx
+
+
+# ===================== Dataset =====================
+
+class DrugPerformerModuleDataset(Dataset):
+    """Drug + PerformerModule Dataset"""
+
+    def __init__(
+        self,
+        drug_data: list,
+        targets: np.ndarray,
+        index_list: list,
+        cell_emb_dict: Dict[str, np.ndarray],
+    ):
+        super().__init__()
+        self.drug_data = drug_data
+        self.targets = targets.astype(np.float32)
+        self.index_list = index_list
+        self.cell_emb_dict = cell_emb_dict
+
+    def __len__(self) -> int:
+        return int(self.targets.shape[0])
+
+    def __getitem__(self, idx: int):
+        # Drug map features
+        feat, adj = self.drug_data[idx]
+        drug_feat = torch.from_numpy(feat).float()
+        drug_adj = torch.from_numpy(adj).float()
+
+        # PerformerModule cell embedding
+        cell_line_id = self.index_list[idx][0]
+        cell_emb_np = self.cell_emb_dict[cell_line_id]
+        cell_emb = torch.from_numpy(cell_emb_np).float()
+
+        # target value
+        target = torch.tensor(self.targets[idx], dtype=torch.float32)
+
+        return {
+            "drug_feat": drug_feat,
+            "drug_adj": drug_adj,
+            "cell_emb": cell_emb,
+            "target": target,
+        }
+
+
+# ===================== Model Definition =====================
+
+class GraphConv(nn.Module):
+    """graph convolution layer"""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter(torch.Tensor(in_features, out_features))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_features))
+        else:
+            self.bias = None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.weight)
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
+        xw = torch.matmul(x, self.weight)
+        if self.bias is not None:
+            xw = xw + self.bias
+        out = torch.bmm(adj.transpose(1, 2), xw)
+        return out
+
+
+class DrugPerformerModuleModel(nn.Module):
+    """
+    Drugs + PerformerModule Predictive Model
+    
+    Architecture:
+    - Drug branch: 3-layer GraphConv + residual connection + double pooling
+    - Cell branch: 2-layer MLP (processing PerformerModule module embeddings)
+    - Fusion layer: 4-layer MLP + attention mechanism
+    """
+
+    def __init__(
+        self,
+        drug_feat_dim: int,
+        cell_emb_dim: int,
+        hidden_dim: int = 256,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(dropout)
+
+        # Drug branch: enhanced 3-layer GraphConv + residual connection
+        self.graph_conv1 = GraphConv(drug_feat_dim, 256)
+        self.graph_conv2 = GraphConv(256, 256)
+        self.graph_conv3 = GraphConv(256, 128)
+        self.drug_bn1 = nn.BatchNorm1d(256)
+        self.drug_bn2 = nn.BatchNorm1d(256)
+        self.drug_bn3 = nn.BatchNorm1d(128)
+
+        # Cell branch: 2-layer MLP (processing PerformerModule embeddings)
+        self.cell_fc1 = nn.Linear(cell_emb_dim, 256)
+        self.cell_bn1 = nn.BatchNorm1d(256)
+        self.cell_fc2 = nn.Linear(256, 128)
+        self.cell_bn2 = nn.BatchNorm1d(128)
+
+        # Fusion dimension: 128 (drug) + 128 (cell) = 256
+        fusion_dim = 256
+
+        # attention mechanism
+        self.attention = nn.Sequential(
+            nn.Linear(fusion_dim, fusion_dim // 2),
+            nn.ReLU(),
+            nn.Linear(fusion_dim // 2, fusion_dim),
+            nn.Sigmoid()
+        )
+
+        # Fusion layer: 4-layer deep MLP
+        self.fusion_fc1 = nn.Linear(fusion_dim, hidden_dim * 2)
+        self.fusion_bn1 = nn.BatchNorm1d(hidden_dim * 2)
+        self.fusion_fc2 = nn.Linear(hidden_dim * 2, hidden_dim)
+        self.fusion_bn2 = nn.BatchNorm1d(hidden_dim)
+        self.fusion_fc3 = nn.Linear(hidden_dim, hidden_dim // 2)
+        self.fusion_bn3 = nn.BatchNorm1d(hidden_dim // 2)
+        self.output_fc = nn.Linear(hidden_dim // 2, 1)
+
+        print_sys(f"Model initialized: drug_feat_dim={drug_feat_dim}, "
+                  f"cell_emb_dim={cell_emb_dim}, hidden_dim={hidden_dim}")
+
+    def forward(
+        self,
+        drug_feat: torch.Tensor,
+        drug_adj: torch.Tensor,
+        cell_emb: torch.Tensor,
+    ) -> torch.Tensor:
+        # Drug branch (enhanced version)
+        # first floor
+        x_drug = self.graph_conv1(drug_feat, drug_adj)  # (B, N, 256)
+        x_drug = x_drug.transpose(1, 2)  # (B, 256, N) for BatchNorm
+        x_drug = self.drug_bn1(x_drug)
+        x_drug = x_drug.transpose(1, 2)  # (B, N, 256)
+        x_drug = self.relu(x_drug)
+        x_drug = self.dropout(x_drug)
+
+        # Second layer + residual connection
+        x_drug_res = x_drug
+        x_drug = self.graph_conv2(x_drug, drug_adj)  # (B, N, 256)
+        x_drug = x_drug.transpose(1, 2)
+        x_drug = self.drug_bn2(x_drug)
+        x_drug = x_drug.transpose(1, 2)
+        x_drug = self.relu(x_drug + x_drug_res)  # residual connection
+        x_drug = self.dropout(x_drug)
+
+        # third floor
+        x_drug = self.graph_conv3(x_drug, drug_adj)  # (B, N, 128)
+        x_drug = x_drug.transpose(1, 2)
+        x_drug = self.drug_bn3(x_drug)
+        x_drug = x_drug.transpose(1, 2)
+        x_drug = self.relu(x_drug)
+
+        # Aggregation: max pooling + mean pooling
+        x_drug_max = x_drug.max(dim=1)[0]  # (B, 128)
+        x_drug_mean = x_drug.mean(dim=1)   # (B, 128)
+        x_drug = x_drug_max + x_drug_mean  # Fusion
+
+        # Cell branching (handling PerformerModule embeddings)
+        x_cell = self.cell_fc1(cell_emb)  # (B, 256)
+        x_cell = self.cell_bn1(x_cell)
+        x_cell = self.relu(x_cell)
+        x_cell = self.dropout(x_cell)
+
+        x_cell = self.cell_fc2(x_cell)  # (B, 128)
+        x_cell = self.cell_bn2(x_cell)
+        x_cell = self.relu(x_cell)
+
+        # Feature fusion + attention
+        h = torch.cat([x_drug, x_cell], dim=1)  # (B, 256)
+
+        # Apply attention mechanism
+        att_weights = self.attention(h)
+        h = h * att_weights  # weighted features
+
+        # Deep Fusion MLP
+        h = self.fusion_fc1(h)  # (B, hidden_dim * 2)
+        h = self.fusion_bn1(h)
+        h = self.relu(h)
+        h = self.dropout(h)
+
+        h = self.fusion_fc2(h)  # (B, hidden_dim)
+        h = self.fusion_bn2(h)
+        h = self.relu(h)
+        h = self.dropout(h)
+
+        h = self.fusion_fc3(h)  # (B, hidden_dim // 2)
+        h = self.fusion_bn3(h)
+        h = self.relu(h)
+        h = self.dropout(h)
+
+        # output layer
+        out = self.output_fc(h)  # (B, 1)
+        return out
+
+
+# ===================== Training and Evaluation =====================
+
+def train_epoch(model, loader, optimizer, criterion, device):
+    """train for one epoch"""
+    model.train()
+    total_loss = 0.0
+    total = 0
+    for batch in loader:
+        drug_feat = batch["drug_feat"].to(device)
+        drug_adj = batch["drug_adj"].to(device)
+        cell_emb = batch["cell_emb"].to(device)
+        target = batch["target"].to(device)
+
+        optimizer.zero_grad(set_to_none=True)
+        pred = model(drug_feat, drug_adj, cell_emb).squeeze(-1)
+
+        loss = criterion(pred, target)
+        loss.backward()
+
+        # gradient crop
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+        optimizer.step()
+
+        total_loss += float(loss.item()) * target.size(0)
+        total += target.size(0)
+    return total_loss / total if total > 0 else 0.0
+
+
+def eval_epoch(model, loader, device):
+    """Evaluate an epoch"""
+    model.eval()
+    preds = []
+    trues = []
+    with torch.no_grad():
+        for batch in loader:
+            drug_feat = batch["drug_feat"].to(device)
+            drug_adj = batch["drug_adj"].to(device)
+            cell_emb = batch["cell_emb"].to(device)
+            target = batch["target"].to(device)
+
+            pred = model(drug_feat, drug_adj, cell_emb).squeeze(-1)
+
+            preds.append(pred.cpu().numpy())
+            trues.append(target.cpu().numpy())
+
+    preds = np.concatenate(preds, axis=0)
+    trues = np.concatenate(trues, axis=0)
+    pcc = pearsonr(preds, trues)[0]
+    return preds, trues, pcc
+
+
+# ===================== Main =====================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Drug response prediction using Drug + PerformerModule"
+    )
+    parser.add_argument("--gpu_id", type=str, default="0")
+    parser.add_argument("--data_path", type=str, default="../drug_response_pytorch/data")
+    parser.add_argument("--epochs", type=int, default=300)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--max_atoms", type=int, default=100)
+    parser.add_argument("--hidden_dim", type=int, default=256)
+    parser.add_argument("--learning_rate", type=float, default=1e-3)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--patience", type=int, default=30)
+
+    # PerformerModule related paths
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        required=True,
+        help="Path to pre-trained PerformerModule checkpoint"
+    )
+    parser.add_argument(
+        "--hyper_params_path",
+        type=str,
+        required=True,
+        help="Path to PerformerModule model hyper parameters pkl file"
+    )
+    parser.add_argument(
+        "--module_data_dir",
+        type=str,
+        default=None,
+        help="Directory containing module JSON files"
+    )
+    parser.add_argument(
+        "--module_data_path",
+        type=str,
+        default=None,
+        help="Path to module JSON file (single file)"
+    )
+    parser.add_argument(
+        "--gene_list_path",
+        type=str,
+        default="../../selected_genes_27k.txt",
+        help="Path to gene list file"
+    )
+    parser.add_argument(
+        "--gene_mapping_path",
+        type=str,
+        default="../../gene2ens.csv",
+        help="Path to gene symbol to Ensembl ID mapping"
+    )
+    parser.add_argument(
+        "--gene2vec_path",
+        type=str,
+        default="../../selected_gene2vec_27k.npy",
+        help="Path to gene2vec file"
+    )
+    parser.add_argument(
+        "--symbol2ens_path",
+        type=str,
+        default="../../gene2ens.csv",
+        help="Path to symbol to Ensembl mapping (for data processing)"
+    )
+    parser.add_argument(
+        "--pool_method",
+        type=str,
+        default="mean",
+        choices=["mean", "max", "cls", "concat"],
+        help="Pooling method for PerformerModule module embeddings"
+    )
+    parser.add_argument(
+        "--shuffle_module_genes",
+        action="store_true",
+        help="Optional downstream-only extra reshuffle on top of module_sets restored from PKL. "
+             "Do not pass this when reproducing the pretrained backbone as-is."
+    )
+    parser.add_argument(
+        "--shuffle_seed",
+        type=int,
+        default=2021,
+        help="Random seed for the downstream extra reshuffle when --shuffle_module_genes is enabled."
+    )
+    # Output path
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="./models_drug_performer_module",
+        help="Directory to save the model"
+    )
+
+    args = parser.parse_args()
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    print_sys("=" * 70)
+    print_sys("Drug Response Prediction: Drug + PerformerModule")
+    print_sys("=" * 70)
+    print_sys(f"Device: {device}")
+    print_sys(f"Epochs: {args.epochs}")
+    print_sys(f"Batch size: {args.batch_size}")
+    print_sys(f"Learning rate: {args.learning_rate}")
+    print_sys(f"Hidden dim: {args.hidden_dim}")
+    print_sys(f"Dropout: {args.dropout}")
+    print_sys(f"Pool method: {args.pool_method}")
+    print_sys(f"Early stopping patience: {args.patience}")
+    print_sys("=" * 70)
+
+    # Create save directory
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    dpath = args.data_path
+    drug_info_file = f"{dpath}/GDSC/1.Drug_listMon Jun 24 09_00_55 2019.csv"
+    cell_line_info_file = f"{dpath}/CCLE/Cell_lines_annotations_20181226.txt"
+    drug_feature_dir = f"{dpath}/GDSC/drug_graph_feat"
+    gene_expression_file = (
+        f"{dpath}/CCLE/genomic_expression_561celllines_697genes_demap_features.csv"
+    )
+
+    # fixed random seed
+    random.seed(0)
+    np.random.seed(0)
+    torch.manual_seed(0)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(0)
+
+    # 1) Read metadata
+    print_sys("\n[Step 1/6] Loading metadata...")
+    drug_feature, gexpr_feature, data_idx = metadata_generate(
+        drug_info_file,
+        cell_line_info_file,
+        drug_feature_dir,
+        gene_expression_file,
+    )
+
+    # 2) Build PerformerModule embeddings
+    print_sys("\n[Step 2/6] Building PerformerModule embeddings...")
+    cell_emb_dict, cell_emb_dim = build_cell_embeddings_with_performer_module( # cell_emb_dict │ Dict[str, np.ndarray] │ {cell_line_id: embedding} mapping dictionary
+        gexpr_feature=gexpr_feature,
+        symbol2ens_path=args.symbol2ens_path,
+        gene_list_path=args.gene_list_path,
+        gene_mapping_path=args.gene_mapping_path,
+        gene2vec_path=args.gene2vec_path,
+        checkpoint_path=args.checkpoint_path,
+        hyper_params_path=args.hyper_params_path,
+        module_data_dir=args.module_data_dir,
+        module_data_path=args.module_data_path,
+        device=device,
+        pool_method=args.pool_method,
+        shuffle_module_genes=args.shuffle_module_genes,
+        seed=args.shuffle_seed,
+    )
+
+    # 3) Divide the data set
+    print_sys("\n[Step 3/6] Splitting dataset...")
+    train_idx, test_idx = data_split(data_idx)
+    print_sys(f"Train samples: {len(train_idx)}, Test samples: {len(test_idx)}")
+
+    # 4) Extract features
+    print_sys("\n[Step 4/6] Extracting features...")
+    X_drug_train, Y_train, train_cancer_info, train_index_list = feature_extract(
+        train_idx, drug_feature, args.max_atoms
+    )
+    X_drug_test, Y_test, test_cancer_info, test_index_list = feature_extract(
+        test_idx, drug_feature, args.max_atoms
+    )
+
+    # 5) Build a data set
+    print_sys("\n[Step 5/6] Creating datasets...")
+    train_ds = DrugPerformerModuleDataset(
+        X_drug_train, Y_train, train_index_list, cell_emb_dict
+    )
+    test_ds = DrugPerformerModuleDataset(
+        X_drug_test, Y_test, test_index_list, cell_emb_dict
+    )
+
+    drug_dim = X_drug_train[0][0].shape[-1]
+
+    num_workers = min(4, os.cpu_count() or 1)
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+    test_loader = DataLoader(
+        test_ds,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+
+    # 6) Build the model
+    print_sys("\n[Step 6/6] Building model...")
+    model = DrugPerformerModuleModel(
+        drug_feat_dim=drug_dim,
+        cell_emb_dim=cell_emb_dim,
+        hidden_dim=args.hidden_dim,
+        dropout=args.dropout,
+    ).to(device)
+
+    # Statistical parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print_sys(f"Total parameters: {total_params:,}")
+    print_sys(f"Trainable parameters: {trainable_params:,}")
+
+    # Optimizers and schedulers
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=args.learning_rate, weight_decay=1e-5
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+        optimizer, T_0=20, T_mult=2, eta_min=1e-6
+    )
+    criterion = torch.nn.HuberLoss(delta=1.0)
+
+    # train
+    print_sys("\n" + "=" * 70)
+    print_sys("Starting training...")
+    print_sys("=" * 70)
+
+    best_pcc = -1.0
+    best_pred = None
+    best_true = None
+    best_epoch = 0
+    patience_counter = 0
+
+    for epoch in range(1, args.epochs + 1):
+        train_loss = train_epoch(model, train_loader, optimizer, criterion, device)
+        preds, trues, pcc = eval_epoch(model, test_loader, device)
+
+        scheduler.step()
+        current_lr = optimizer.param_groups[0]["lr"]
+
+        print_sys(
+            f"Epoch {epoch:03d}/{args.epochs} | "
+            f"train_loss={train_loss:.4f} | "
+            f"val_pcc={pcc:.4f} | "
+            f"lr={current_lr:.6f}"
+        )
+
+        if pcc > best_pcc:
+            best_pcc = pcc
+            best_pred = preds
+            best_true = trues
+            best_epoch = epoch
+            patience_counter = 0
+
+            # Save the best model
+            model_save_path = os.path.join(args.save_dir, "best_model.pth")
+            torch.save(
+                {
+                    "epoch": epoch,
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "scheduler_state_dict": scheduler.state_dict(),
+                    "best_pcc": best_pcc,
+                    "train_loss": train_loss,
+                    "args": vars(args),
+                },
+                model_save_path,
+            )
+            print_sys(f"  → Best model saved! PCC improved to {best_pcc:.4f}")
+        else:
+            patience_counter += 1
+            if patience_counter >= args.patience:
+                print_sys(
+                    f"\n[Early Stopping] No improvement for {args.patience} epochs."
+                )
+                break
+
+    # Save final result
+    print_sys("\n" + "=" * 70)
+    print_sys("Training completed!")
+    print_sys("=" * 70)
+    print_sys(f"Best epoch: {best_epoch}")
+    print_sys(f"Best PCC: {best_pcc:.4f}")
+
+    if best_pred is not None:
+        result_path = os.path.join(args.save_dir, "predictions.npz")
+        np.savez(result_path, pred=best_pred, test=best_true)
+        print_sys(f"Predictions saved to: {result_path}")
+
+    print_sys(f"Best model saved to: {os.path.join(args.save_dir, 'best_model.pth')}")
+    print_sys("=" * 70)
+
+
+if __name__ == "__main__":
+    main()
+

BioMod/environment.yml

@@ -0,0 +1,68 @@
+name: performer_module
+channels:
+  - pytorch
+  - nvidia
+  - pyg
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.9
+  - pip>=24.0
+
+  # GPU stack matching the current server environment.
+  # For a different CUDA driver/toolkit, adjust pytorch-cuda and the PyG notes
+  # in requirements.txt accordingly.
+  - pytorch=2.6.*
+  - torchvision=0.21.*
+  - torchaudio=2.6.*
+  - pytorch-cuda=12.4
+  - pyg=2.6.*
+
+  # Core scientific stack
+  - numpy=1.26.*
+  - scipy=1.13.*
+  - pandas>=2.2,<2.4
+  - scikit-learn=1.6.*
+  - h5py>=3.10
+  - anndata=0.10.*
+  - scanpy=1.10.*
+  - matplotlib=3.9.*
+  - seaborn=0.13.*
+  - tqdm>=4.66
+  - requests>=2.31
+  - pyyaml>=6.0
+  - networkx>=3.2
+  - statsmodels>=0.14
+  - umap-learn>=0.5
+  - numba
+  - openpyxl
+  - jupyterlab
+  - ipykernel
+
+  # Single-cell and graph-analysis helpers
+  - python-igraph>=0.11
+  - leidenalg>=0.10
+  - rdkit
+
+  - pip:
+      # PerformerModule model utilities
+      - einops==0.8.1
+      - local-attention==1.11.2
+      - hickle>=5.0.0
+      - dcor==0.6
+
+      # Optional task modules used by downstream scripts.
+      - scib>=1.1.5
+      - decoupler>=1.8.0
+      - gseapy>=1.1.0
+      - mygene>=3.2.2
+      - lifelines>=0.27.0
+
+      # Text-alignment and LLM-assisted annotation utilities.
+      - sentence-transformers>=2.7.0
+      - transformers>=4.40.0
+      - openai>=1.40.0
+      - wandb==0.23.0
+
+      # Only needed if running legacy pickles on Python < 3.8.
+      - pickle5>=0.0.12; python_version < "3.8"

BioMod/eval_new_data/cmd.txtx

@@ -0,0 +1,19 @@
+# masked 模式（与预训练 valid loss 直接可比）
+python evaluate_new_data.py \
+--ckpt_dir ../mae_tost_48M_cross_attn_only_pure_memory/ \
+--model_name mae_tost_48M_cross_attn_only_pure_memory \
+--data_dir ./npz \
+--module_data_dir ../module_data_all/ \
+--eval_mode masked \
+--batch_size 320 \
+--output_csv eval_masked.csv 
+
+# full_recon 模式（不 mask，全 module 重建 loss） 
+python evaluate_new_data.py \
+--ckpt_dir /path/to/mae_tost_48M_cross_attn_only_pure_memory/ \
+--model_name mae_tost_48M_cross_attn_only_pure_memory \
+--data_dir /path/to/20260423测试集/npz/npz \ 
+--module_data_dir ./module_data_all/ \ 
+--eval_mode full_recon \
+--batch_size 32 \
+--output_csv eval_full_recon.csv 

BioMod/eval_new_data/evaluate_new_data.py

@@ -0,0 +1,375 @@
+# -*- coding: utf-8 -*-
+"""
+Evaluate the reconstruction loss of a pre-trained model on a new dataset (DDP multi-card version).
+Used to verify the generalization ability of the model on unseen donor data, compared with pre-training valid loss.
+
+Two evaluation modes are supported:
+  masked — consistent with pre-training validation, using the same mask_ratio, for direct comparison
+  full_recon — without masking, calculate the reconstruction loss of all modules
+
+Usage example (single card):
+  torchrun --nproc_per_node=1 evaluate_new_data.py \
+      --ckpt_dir ./ckpts/ \
+      --model_name mae_tost_48M_cross_attn_only_pure_memory \
+      --data_dir /path/to/20260423test set/npz/npz \
+      --eval_mode masked --batch_size 32
+
+Usage example (4 card DDP):
+  torchrun --nproc_per_node=4 evaluate_new_data.py \
+      --ckpt_dir ./ckpts/ \
+      --model_name mae_tost_48M_cross_attn_only_pure_memory \
+      --data_dir /path/to/20260423test set/npz/npz \
+      --eval_mode masked --batch_size 32
+"""
+
+import os
+import re
+import csv
+import sys
+import glob
+import time
+import pickle
+import inspect
+import argparse
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from scipy import sparse
+
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+
+from models_performer_module import PerformerEncoderOnly_Module
+from module_dataset import load_module_data
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+class DonorSCDataset(Dataset):
+    """Dataset of a single donor, returning single-cell expression vectors row-wise from dense numpy."""
+
+    def __init__(self, csr_matrix):
+        self.data = csr_matrix.toarray()
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    def __getitem__(self, idx):
+        return torch.from_numpy(self.data[idx].astype(np.float32))
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint discovery
+# ---------------------------------------------------------------------------
+
+def discover_checkpoints(ckpt_dir, model_name, epoch_filter=None):
+    pattern = os.path.join(ckpt_dir, f"{model_name}_epoch*.pth")
+    files = glob.glob(pattern)
+    results = []
+    for f in files:
+        match = re.search(r'epoch(\d+)', os.path.basename(f))
+        if match:
+            epoch_num = int(match.group(1))
+            if epoch_filter is None or epoch_num in epoch_filter:
+                results.append((epoch_num, f))
+    results.sort(key=lambda x: x[0])
+    return results
+
+
+def load_model_config(ckpt_dir, model_name):
+    pkl_path = os.path.join(ckpt_dir, f"{model_name}.pkl")
+    if not os.path.isfile(pkl_path):
+        raise FileNotFoundError(f"Model configuration file not found: {pkl_path}")
+    with open(pkl_path, 'rb') as f:
+        return pickle.load(f)
+
+
+def filter_model_config(config):
+    """Filter out obsolete parameters in pkl that are not accepted by the model __init__."""
+    valid = set(inspect.signature(PerformerEncoderOnly_Module.__init__).parameters.keys())
+    removed = {k: config.pop(k) for k in list(config) if k not in valid and k != "self"}
+    return removed
+
+
+# ---------------------------------------------------------------------------
+# Evaluation
+# ---------------------------------------------------------------------------
+
+def evaluate_donor_ddp(model, dataloader, device, eval_mode, mask_ratio, ptdtype, seed=None):
+    """
+    Calculate the loss for this rank shard of a single donor and return (loss_sum, batch_count).
+    Aggregation is done by the caller via all_reduce.
+    """
+    model.eval()
+    loss_sum = 0.0
+    batch_count = 0
+
+    raw_model = model.module if isinstance(model, DDP) else model
+
+    with torch.no_grad():
+        for batch_idx, data in enumerate(dataloader):
+            data = data.to(device, non_blocking=True)
+            if torch.isnan(data).any() or torch.isinf(data).any():
+                continue
+
+            if eval_mode == "masked" and seed is not None:
+                torch.manual_seed(seed + batch_idx)
+                torch.cuda.manual_seed(seed + batch_idx)
+
+            with torch.autocast(device_type='cuda', dtype=ptdtype):
+                if eval_mode == "masked":
+                    loss, _, _ = raw_model(data, mask_ratio=mask_ratio)
+                else:
+                    latent, mask_zero, ids_restore, module_targets = raw_model.forward_encoder(
+                        data, mask_ratio=0.0
+                    )
+                    dt = raw_model.decoder_type
+                    if dt == "performer":
+                        preds = raw_model.forward_decoder_performer(latent, ids_restore)
+                    elif dt == "cross_attn":
+                        preds = raw_model.forward_decoder_cross_attn(latent, mask_zero, ids_restore)
+                    elif dt == "simple":
+                        preds = raw_model.forward_decoder_simple(latent, ids_restore)
+                    elif dt == "none":
+                        preds = raw_model.forward_predictor_from_encoder(latent)
+                    else:
+                        raise RuntimeError(f"Unsupported decoder_type: {dt}")
+                    mask_all = torch.ones_like(mask_zero)
+                    loss = raw_model.compute_loss(preds, module_targets, mask_all)
+
+            if torch.isnan(loss) or torch.isinf(loss):
+                continue
+            loss_sum += loss.item()
+            batch_count += 1
+
+    return loss_sum, batch_count
+
+
+def reduce_scalar(value, device):
+    """all_reduce is a scalar and returns the global summation result."""
+    t = torch.tensor(value, dtype=torch.float64, device=device)
+    dist.all_reduce(t, op=dist.ReduceOp.SUM)
+    return t.item()
+
+
+def discover_datasets(data_dir):
+    datasets = []
+    for name in sorted(os.listdir(data_dir)):
+        subdir = os.path.join(data_dir, name)
+        if os.path.isdir(subdir):
+            npz_files = sorted([f for f in os.listdir(subdir) if f.endswith('.npz')])
+            if npz_files:
+                datasets.append((name, subdir, npz_files))
+    return datasets
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate pre-trained model reconstruction loss (DDP) on new datasets")
+    parser.add_argument("--ckpt_dir", type=str, required=True)
+    parser.add_argument("--model_name", type=str, required=True)
+    parser.add_argument("--data_dir", type=str, required=True)
+    parser.add_argument("--gene_list_path", type=str, default="../../selected_genes_27k.txt")
+    parser.add_argument("--gene_mapping_path", type=str, default="../../gene2ens.csv")
+    parser.add_argument("--module_data_dir", type=str, default=None)
+    parser.add_argument("--gene2vec_data_path", type=str, default="../../selected_gene2vec_27k.npy")
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--output_csv", type=str, default="eval_new_data_results.csv")
+    parser.add_argument("--epochs", type=str, default=None,
+                        help="Comma separated list of epochs (e.g. '1,5,10'), default is all")
+    parser.add_argument("--eval_mode", type=str, default="masked",
+                        choices=["masked", "full_recon"])
+    parser.add_argument("--mask_ratio", type=float, default=0.75)
+    parser.add_argument("--seed", type=int, default=2021)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--amp_dtype", type=str, default="bfloat16",
+                        choices=["float32", "bfloat16", "float16"])
+    args = parser.parse_args()
+
+    # ---- DDP initialization ----
+    dist.init_process_group(backend='nccl')
+    local_rank = int(os.environ.get('LOCAL_RANK', 0))
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    device = torch.device(f'cuda:{local_rank}')
+    torch.cuda.set_device(device)
+    is_master = (rank == 0)
+
+    ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.amp_dtype]
+
+    if is_master:
+        print(f"World size: {world_size}", flush=True)
+        print(f"Eval mode: {args.eval_mode}", flush=True)
+        if args.eval_mode == "masked":
+            print(f"Mask ratio: {args.mask_ratio}", flush=True)
+
+    # ---- Load model configuration ----
+    if is_master:
+        print(f"\nLoad model configuration: {args.ckpt_dir}/{args.model_name}.pkl", flush=True)
+    model_config = load_model_config(args.ckpt_dir, args.model_name)
+    removed = filter_model_config(model_config)
+    if is_master and removed:
+        for k, v in removed.items():
+            print(f"  Skip obsolete parameters: {k}={v!r}", flush=True)
+
+    if args.gene2vec_data_path:
+        model_config["gene2vec_file"] = args.gene2vec_data_path
+
+    if "module_sets" in model_config and model_config["module_sets"]:
+        if is_master:
+            print(f"Loaded from pkl config {len(model_config['module_sets'])} modules", flush=True)
+    else:
+        if args.module_data_dir is None:
+            raise ValueError("There are no module_sets in the model configuration, please specify via --module_data_dir")
+        model_config["module_sets"] = load_module_data(
+            module_data_dir=args.module_data_dir,
+            gene_list_path=args.gene_list_path,
+            gene_mapping_path=args.gene_mapping_path,
+            verbose=is_master,
+        )
+
+    # ---- Discover checkpoints & data sets ----
+    epoch_filter = set(int(e) for e in args.epochs.split(",")) if args.epochs else None
+    checkpoints = discover_checkpoints(args.ckpt_dir, args.model_name, epoch_filter)
+    datasets = discover_datasets(args.data_dir)
+
+    if is_master:
+        if not checkpoints:
+            print(f"checkpoint not found: {args.ckpt_dir}/{args.model_name}_epoch*.pth", flush=True)
+        else:
+            print(f"turn up {len(checkpoints)} checkpoints: {[e for e, _ in checkpoints]}", flush=True)
+        for ds_name, _, npz_files in datasets:
+            print(f"Dataset {ds_name}: {len(npz_files)} donors", flush=True)
+        if not datasets:
+            print(f"Warning: in {args.data_dir} The subdirectory containing the .npz file was not found!", flush=True)
+
+    if not checkpoints or not datasets:
+        dist.destroy_process_group()
+        return
+
+    # ----CSV (master write only)----
+    csv_file, writer = None, None
+    if is_master:
+        csv_file = open(args.output_csv, 'w', newline='', encoding='utf-8')
+        writer = csv.writer(csv_file)
+        writer.writerow(["epoch", "dataset", "donor", "num_cells", "mean_loss", "eval_mode"])
+
+    # ----Build the model (create it only once, only reload the weights every epoch)----
+    model = PerformerEncoderOnly_Module(**model_config)
+    model.to(device)
+
+    # ---- Main loop ----
+    summary = defaultdict(float)
+
+    for epoch_num, ckpt_path in checkpoints:
+        if is_master:
+            print(f"\n{'='*60}", flush=True)
+            print(f"Epoch {epoch_num}: {os.path.basename(ckpt_path)}", flush=True)
+            print(f"{'='*60}", flush=True)
+
+        ckpt = torch.load(ckpt_path, map_location=device)
+        model.load_state_dict(ckpt['model_state_dict'])
+        del ckpt
+        torch.cuda.empty_cache()
+        model.eval()
+
+        for ds_name, ds_dir, npz_files in datasets:
+            ds_losses = []
+            ds_cells_total = 0
+            t_ds_start = time.time()
+
+            for donor_idx, npz_file in enumerate(npz_files):
+                donor_id = npz_file.replace('.npz', '')
+                npz_path = os.path.join(ds_dir, npz_file)
+
+                data_npz = np.load(npz_path, allow_pickle=True)
+                csr = sparse.csr_matrix(
+                    (data_npz['data'], data_npz['indices'], data_npz['indptr']),
+                    shape=tuple(data_npz['shape'])
+                )
+                num_cells = csr.shape[0]
+
+                dataset_obj = DonorSCDataset(csr)
+                sampler = DistributedSampler(dataset_obj, num_replicas=world_size,
+                                             rank=rank, shuffle=False, drop_last=False)
+                loader = DataLoader(
+                    dataset_obj, batch_size=args.batch_size, sampler=sampler,
+                    num_workers=args.num_workers, pin_memory=True, drop_last=False,
+                )
+
+                donor_seed = args.seed + donor_idx if args.eval_mode == "masked" else None
+                local_loss_sum, local_batch_count = evaluate_donor_ddp(
+                    model, loader, device, args.eval_mode,
+                    args.mask_ratio, ptdtype, seed=donor_seed,
+                )
+
+                global_loss_sum = reduce_scalar(local_loss_sum, device)
+                global_batch_count = reduce_scalar(local_batch_count, device)
+
+                if global_batch_count > 0:
+                    mean_loss = global_loss_sum / global_batch_count
+                else:
+                    mean_loss = 0.0
+
+                if is_master:
+                    writer.writerow([epoch_num, ds_name, donor_id, num_cells,
+                                     f"{mean_loss:.6f}", args.eval_mode])
+                    csv_file.flush()
+                    ds_losses.append(mean_loss)
+                    ds_cells_total += num_cells
+
+                    if (donor_idx + 1) % 10 == 0 or (donor_idx + 1) == len(npz_files):
+                        print(f"  {ds_name}: {donor_idx+1}/{len(npz_files)} donors done", flush=True)
+
+                dist.barrier()
+
+            if is_master:
+                ds_mean = np.mean(ds_losses) if ds_losses else 0.0
+                print(f"  {ds_name}: mean_loss={ds_mean:.6f}, "
+                      f"total_cells={ds_cells_total}, time={time.time()-t_ds_start:.1f}s", flush=True)
+                summary[(epoch_num, ds_name)] = ds_mean
+
+        torch.cuda.empty_cache()
+        dist.barrier()
+
+    del model
+    torch.cuda.empty_cache()
+
+    # ---- Summary table ----
+    if is_master:
+        csv_file.close()
+        print(f"\nDetailed results have been saved to: {args.output_csv}")
+        print(f"\n{'='*60}")
+        print("summary table")
+        print(f"{'='*60}")
+        ds_names = sorted(set(ds for _, ds in summary.keys()))
+        header = f"{'Epoch':>8}" + "".join(f"{ds:>16}" for ds in ds_names)
+        print(header)
+        print("-" * len(header))
+        for epoch_num, _ in checkpoints:
+            row = f"{epoch_num:>8}"
+            for ds in ds_names:
+                row += f"{summary.get((epoch_num, ds), float('nan')):>16.6f}"
+            print(row)
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    import traceback
+    try:
+        main()
+    except Exception:
+        traceback.print_exc()
+        sys.stdout.flush()
+        sys.stderr.flush()
+        sys.exit(1)

BioMod/gene_inverse_norman/README.md

@@ -0,0 +1,217 @@
+# Norman Reverse Perturbation Prediction (scGPT / CellFM Protocol Alignment)
+
+This directory implements reverse perturbation prediction:
+
+- Input: single-cell expression after perturbation.
+- Candidate space: all single and double combinations from the selected gene pool; 20 genes produce 210 candidates.
+- Training objective: identify the true perturbation within the candidate universe.
+- Split strategy: split by perturbation case, not by cell.
+- Main evaluation metrics: `condition_top1_acc` to `condition_top10_acc` for exact hits, plus `condition_relevant_top1_acc` to `condition_relevant_top10_acc` for predictions that hit at least one true gene.
+- Paper-style supplementary metrics: `case_cellmean_top1_acc` to `case_cellmean_top10_acc`, which first compute cell-level hit rates within each condition and then average conditions equally; `case_cellmean_relevant_topk_acc` is also supported.
+
+Two built-in protocols are supported and can be selected with `--protocol scgpt`
+or `--protocol cellfm`.
+
+## Protocol Alignment
+
+Both protocols come from the exact condition lists reported in the papers.
+Canonicalization maps `GENE+GENE`, `ctrl+GENE`, and `GENE+ctrl` single-gene
+forms to `GENE`, matching the Norman annotations.
+
+| Item | scGPT | CellFM |
+| --- | --- | --- |
+| Gene pool | 20 | 20 |
+| Candidate universe | 210 (20 singles + 190 doubles) | 210 (20 singles + 190 doubles) |
+| Train cases | 39 (20 doubles + 19 singles) | 33 (16 doubles + 17 singles) |
+| Val cases | 3 doubles | 3 doubles |
+| Test cases | 7 doubles | 11 doubles |
+| Total observed cases used | 49 | 47 |
+| Missing in local Norman data | 0 | 0 |
+
+In the scGPT gene pool, FOSB is not included as a single-gene training case. In
+the CellFM gene pool, FOSB, FOXA3, and MAP2K3 also have no single-gene training
+entries. These singles exist in Norman and remain in the candidate universe as
+distractors, matching the known-vs-unseen design in the papers.
+
+The exact gene and split lists are stored in:
+
+- `protocols/scgpt/genes.txt` and `protocols/scgpt/split.json`
+- `protocols/cellfm/genes.txt` and `protocols/cellfm/split.json`
+
+## 1. Build the Benchmark
+
+```bash
+# scGPT
+python performer_module/gene_inverse_norman/build_norman_gene_inverse_benchmark.py \
+  --adata_path GEARS/data/norman/perturb_processed.h5ad \
+  --protocol scgpt \
+  --output_dir performer_module/gene_inverse_norman/outputs/bench_scgpt
+
+# CellFM
+python performer_module/gene_inverse_norman/build_norman_gene_inverse_benchmark.py \
+  --adata_path GEARS/data/norman/perturb_processed.h5ad \
+  --protocol cellfm \
+  --output_dir performer_module/gene_inverse_norman/outputs/bench_cellfm
+```
+
+Main outputs:
+
+- `gene_inverse_conditions.csv`: full candidate universe with `is_observed`.
+- `class_mapping.csv`: `class_id -> condition`.
+- `benchmark_summary.json`: protocol, observed counts, and `protocol_conditions_missing_in_data`.
+- `protocol_split.json`: backup copy of the original protocol split.
+
+## 2. Training
+
+```bash
+# scGPT
+python finetune_gene_inverse.py \
+  --scfm_hyper_params_path ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+  --scfm_ckpt_path ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+  --benchmark_csv ./outputs/bench_scgpt/gene_inverse_conditions.csv \
+  --adata_path ../../GEARS/data/norman/perturb_processed.h5ad \
+  --gene_mapping_path ../../gene2ens.csv \
+  --scfm_genes_file ../../selected_genes_27k.txt \
+  --gene2vec_path ../../selected_gene2vec_27k.npy \
+  --protocol scgpt \
+  --finetune_mode full \
+  --pool_method max \
+  --epochs 20 \
+  --batch_size 24 \
+  --lr_head 1e-3 \
+  --lr_backbone 1e-4 \
+  --algorithm_name scLong_scgpt_protocol \
+  --export_test_logits \
+  --save_best_checkpoint \
+  --output_dir performer_module/gene_inverse_norman/outputs/run_scgpt
+
+# CellFM
+python performer_module/gene_inverse_norman/finetune_gene_inverse.py \
+  --scfm_hyper_params_path mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+  --scfm_ckpt_path mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pth \
+  --benchmark_csv performer_module/gene_inverse_norman/outputs/bench_cellfm/gene_inverse_conditions.csv \
+  --adata_path GEARS/data/norman/perturb_processed.h5ad \
+  --gene_mapping_path gene2ens.csv \
+  --scfm_genes_file selected_genes_27k.txt \
+  --gene2vec_path selected_gene2vec_27k.npy \
+  --protocol cellfm \
+  --finetune_mode frozen \
+  --pool_method max \
+  --epochs 5 \
+  --batch_size 4 \
+  --lr_head 1e-4 \
+  --lr_backbone 1e-5 \
+  --algorithm_name scLong_cellfm_protocol \
+  --export_test_logits \
+  --save_best_checkpoint \
+  --output_dir performer_module/gene_inverse_norman/outputs/run_cellfm
+```
+
+Once `--protocol` is specified, the fixed train/val/test lists from
+`protocols/<name>/split.json` are used. Random splitting is disabled. `--seed`
+only affects initialization and dataloader order.
+
+DDP training on multiple GPUs:
+
+```bash
+torchrun --nproc_per_node=4 \
+  performer_module/gene_inverse_norman/finetune_gene_inverse.py \
+  --protocol cellfm \
+  ... # Other arguments are the same as above.
+```
+
+Training outputs:
+
+- `best_finetune_gene_inverse.pth`, saved only when `--save_best_checkpoint` is passed.
+- `summary_metrics.json`, including `test_condition_top1_acc` to `test_condition_top10_acc`, `test_condition_relevant_top1_acc` to `test_condition_relevant_top10_acc`, and `case_cellmean` metrics such as `test_case_cellmean_top1_acc` and `test_case_cellmean_relevant_top1_acc`.
+- `test_summary_metrics.json`, which keeps only the test-set view and key configuration for reporting, without all `best_val_*` fields.
+- `metrics_table.csv`, a flattened top1-to-top10 table grouped by metric family and including the `algorithm` field.
+- `train_history.csv`.
+- `test_predictions.csv` for per-cell predictions.
+- `test_logits.csv` when `--export_test_logits` is enabled.
+- `test_condition_rankings.jsonl`, containing top-10 rankings and the true rank for each test perturbation.
+- `condition_splits.npz`, `condition_splits.json`, and `split_indices.npz`.
+
+Notes:
+
+- Full finetuned checkpoints are not saved by default to save disk space. Add `--save_best_checkpoint` if later re-evaluation with `evaluate_gene_inverse.py` is needed.
+- `--select_metric` is used only to choose the best epoch by one metric; it is not a multi-metric selector. Other metrics are still fully exported.
+- When `--select_metric` ties, `condition_top1_acc`, `case_cellmean_top5_acc`, and `val_loss` are used as tie-breakers.
+- Optional training stabilization: `--class_balance_power 0.5` for inverse-square-root weighting by condition sample count, and `--label_smoothing 0.05`.
+
+## 3. Re-Evaluation
+
+```bash
+python performer_module/gene_inverse_norman/evaluate_gene_inverse.py \
+  --trained_head_ckpt performer_module/gene_inverse_norman/outputs/run_cellfm/best_finetune_gene_inverse.pth \
+  --scfm_hyper_params_path mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+  --scfm_ckpt_path mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pth \
+  --benchmark_csv performer_module/gene_inverse_norman/outputs/bench_cellfm/gene_inverse_conditions.csv \
+  --condition_splits_npz performer_module/gene_inverse_norman/outputs/run_cellfm/condition_splits.npz \
+  --split test \
+  --export_logits \
+  --output_dir performer_module/gene_inverse_norman/outputs/eval_cellfm
+```
+
+## 4. Evaluation Definition
+
+### Main Metrics Aligned with the Papers
+
+`condition_top1_acc` ... `condition_top10_acc`: exact-hit metrics.
+
+`condition_relevant_top1_acc` ... `condition_relevant_top10_acc`: relevant-hit
+metrics, where the predicted set shares at least one gene with the true
+perturbation.
+
+`case_cellmean_top1_acc` ... `case_cellmean_top10_acc`: first compute cell-level
+hit@k within each condition, then average conditions equally. This can produce
+decimal values that are not simply `x/7` or `x/11`.
+
+`case_cellmean_relevant_top1_acc` ... `case_cellmean_relevant_top10_acc`: same
+idea as above, but using the relevant-hit definition.
+
+Workflow:
+
+1. The model outputs logits over 210 candidates for each test cell.
+2. Logits from all cells under the same perturbation are averaged.
+3. The true perturbation is checked against the top-k ranking.
+
+### Auxiliary Metrics
+
+`top1_acc` ... `top10_acc`: cell-level top-k metrics for reference only.
+
+`--select_metric` can be one of `condition_top{k}_acc`,
+`condition_relevant_top{k}_acc`, `case_cellmean_top{k}_acc`,
+`case_cellmean_relevant_top{k}_acc`, or `top{k}_acc`, where `k=1..10`.
+
+Recommendations:
+
+- Use `condition_topk_acc` and `condition_relevant_topk_acc` for the main conclusion.
+- Also report `case_cellmean_*` when aligning with paper wording about non-integer average hit counts.
+- The classification head uses residual connections by default; add `--no_head_residual` to return to the old head.
+- Use only `max` or `mean` for `pool_method`; other historical options have been removed from the current training entry point.
+
+## 5. Points to Watch
+
+- The candidate space is the full set of 210 singles and doubles. Paper-style Top-k hit rates are ranked over this full set by default. If a baseline ranks only a subset, report that separately because the numbers are not directly comparable.
+- The two protocols use different gene pools: scGPT has 19 unique genes plus FEV, while CellFM has 20 unique genes. Do not directly cross-compare scGPT-train and CellFM-test without explaining this.
+- `seed` does not affect the split because splits come from fixed protocol lists, but it still affects initialization and dataloader order. Multi-seed runs can estimate variance.
+
+## 6. Custom Splits Without a Protocol
+
+Custom gene pools and random splits are still supported:
+
+```bash
+python performer_module/gene_inverse_norman/build_norman_gene_inverse_benchmark.py \
+  --adata_path GEARS/data/norman/perturb_processed.h5ad \
+  --gene_pool_file my_genes.txt \
+  --output_dir performer_module/gene_inverse_norman/outputs/bench_custom
+
+python performer_module/gene_inverse_norman/finetune_gene_inverse.py \
+  ... \
+  --train_cases 33 --val_cases 3 --test_cases 11 \
+  --allow_observed_count_mismatch
+```
+
+When `--protocol` is not passed, observed conditions are randomly split by
+`--train_cases`, `--val_cases`, `--test_cases`, and `--seed`.

BioMod/gene_inverse_norman/build_norman_gene_inverse_benchmark.py

@@ -0,0 +1,300 @@
+import argparse
+import itertools
+import json
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+
+CURRENT_DIR = Path(__file__).resolve().parent
+if str(CURRENT_DIR) not in sys.path:
+    sys.path.insert(0, str(CURRENT_DIR))
+
+from utils import canonical_condition_name, default_repo_path, parse_condition, print_sys
+
+
+PROTOCOLS_DIR = CURRENT_DIR / "protocols"
+SUPPORTED_PROTOCOLS = ("scgpt", "cellfm")
+
+
+def protocol_dir(protocol: str) -> Path:
+    if protocol not in SUPPORTED_PROTOCOLS:
+        raise ValueError(f"Unsupported protocol '{protocol}'. Expected one of {SUPPORTED_PROTOCOLS}.")
+    return PROTOCOLS_DIR / protocol
+
+
+def load_protocol_genes(protocol: str) -> List[str]:
+    genes_path = protocol_dir(protocol) / "genes.txt"
+    with open(genes_path, "r", encoding="utf-8") as handle:
+        genes = [line.strip() for line in handle.readlines() if line.strip()]
+    return sorted(dict.fromkeys(genes))
+
+
+def load_protocol_split(protocol: str) -> Dict[str, object]:
+    split_path = protocol_dir(protocol) / "split.json"
+    with open(split_path, "r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def load_gene_pool(args) -> List[str]:
+    if args.protocol:
+        return load_protocol_genes(args.protocol)
+    if args.gene_pool_file:
+        with open(args.gene_pool_file, "r", encoding="utf-8") as handle:
+            genes = [line.strip() for line in handle.readlines() if line.strip()]
+        return sorted(dict.fromkeys(genes))
+    raise ValueError("Must specify --protocol or --gene_pool_file")
+
+
+def build_candidate_conditions(genes: List[str], include_singles: bool = True, include_doubles: bool = True) -> List[str]:
+    candidates: List[str] = []
+    if include_singles:
+        candidates.extend(genes)
+    if include_doubles:
+        candidates.extend([f"{a}+{b}" for a, b in itertools.combinations(genes, 2)])
+    return sorted(candidates)
+
+
+def summarize_observed_conditions(
+    adata,
+    condition_col: str,
+    condition_name_col: Optional[str],
+    candidate_set: set,
+) -> Dict[str, Dict[str, object]]:
+    if condition_col not in adata.obs.columns:
+        raise KeyError(f"condition_col '{condition_col}' not found in adata.obs")
+
+    condition_series = adata.obs[condition_col].astype(str)
+    raw_condition_counts = condition_series.value_counts().to_dict()
+
+    raw_condition_name_map: Dict[str, str] = {}
+    if condition_name_col and condition_name_col in adata.obs.columns:
+        raw_condition_name_map = (
+            adata.obs[[condition_col, condition_name_col]]
+            .astype(str)
+            .drop_duplicates(subset=[condition_col])
+            .set_index(condition_col)[condition_name_col]
+            .to_dict()
+        )
+
+    candidate_to_raw: Dict[str, List[str]] = {candidate: [] for candidate in candidate_set}
+    for raw_condition in sorted(raw_condition_counts.keys()):
+        if raw_condition == "ctrl":
+            continue
+        canonical = canonical_condition_name(raw_condition)
+        if canonical in candidate_set:
+            candidate_to_raw[canonical].append(raw_condition)
+
+    candidate_summary: Dict[str, Dict[str, object]] = {}
+    for candidate in sorted(candidate_set):
+        raw_conditions = candidate_to_raw.get(candidate, [])
+        n_cells = int(sum(raw_condition_counts[raw] for raw in raw_conditions)) if raw_conditions else 0
+        candidate_summary[candidate] = {
+            "raw_conditions": raw_conditions,
+            "condition_names": [raw_condition_name_map.get(raw, "") for raw in raw_conditions],
+            "n_cells": n_cells,
+            "is_observed": int(n_cells > 0),
+        }
+    return candidate_summary
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Build a CellFM-style Norman reverse-perturbation benchmark table."
+    )
+    parser.add_argument(
+        "--adata_path",
+        type=str,
+        default=default_repo_path("GEARS", "data", "norman", "perturb_processed.h5ad"),
+        help="Path to Norman perturb_processed.h5ad",
+    )
+    parser.add_argument("--condition_col", type=str, default="condition", help="Condition column in adata.obs")
+    parser.add_argument(
+        "--condition_name_col",
+        type=str,
+        default="condition_name",
+        help="Optional condition-name column in adata.obs",
+    )
+    parser.add_argument(
+        "--protocol",
+        type=str,
+        default="",
+        choices=["", "scgpt", "cellfm"],
+        help="Use a built-in protocol (scgpt or cellfm) for gene pool and split. Overrides --gene_pool_file.",
+    )
+    parser.add_argument(
+        "--gene_pool_file",
+        type=str,
+        default="",
+        help="Optional newline-delimited gene pool. Used only when --protocol is empty.",
+    )
+    parser.add_argument(
+        "--include_singles",
+        action="store_true",
+        help="Include single-gene candidates. Defaults follow the protocol when --protocol is set.",
+    )
+    parser.add_argument(
+        "--include_doubles",
+        action="store_true",
+        help="Include double-gene candidates. Defaults follow the protocol when --protocol is set.",
+    )
+    parser.add_argument("--singles_only", action="store_true", help="Only keep single-gene candidates.")
+    parser.add_argument("--doubles_only", action="store_true", help="Only keep double-gene candidates.")
+    parser.add_argument(
+        "--min_cells_per_condition",
+        type=int,
+        default=1,
+        help="Observed candidates with fewer cells than this threshold are downgraded to unobserved.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=default_repo_path(
+            "performer_module",
+            "gene_inverse_norman",
+            "outputs",
+            "norman_gene_inverse_benchmark",
+        ),
+        help="Output directory",
+    )
+    args = parser.parse_args()
+
+    if args.singles_only and args.doubles_only:
+        raise ValueError("--singles_only and --doubles_only cannot be used together")
+
+    protocol_split: Optional[Dict[str, object]] = None
+    if args.protocol:
+        protocol_split = load_protocol_split(args.protocol)
+        include_singles = bool(protocol_split.get("include_singles", False))
+        include_doubles = bool(protocol_split.get("include_doubles", True))
+    else:
+        include_singles = True
+        include_doubles = True
+
+    if args.singles_only:
+        include_doubles = False
+    if args.doubles_only:
+        include_singles = False
+    if args.include_singles:
+        include_singles = True
+    if args.include_doubles:
+        include_doubles = True
+    if not include_singles and not include_doubles:
+        raise ValueError("Candidate space cannot be empty")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    selected_genes = load_gene_pool(args)
+    candidate_conditions = build_candidate_conditions(
+        genes=selected_genes,
+        include_singles=include_singles,
+        include_doubles=include_doubles,
+    )
+    candidate_set = set(candidate_conditions)
+
+    print_sys(f"Loading Norman data from {args.adata_path}")
+    adata = ad.read_h5ad(args.adata_path, backed="r")
+
+    candidate_summary = summarize_observed_conditions(
+        adata=adata,
+        condition_col=args.condition_col,
+        condition_name_col=args.condition_name_col,
+        candidate_set=candidate_set,
+    )
+
+    rows = []
+    min_cells = int(args.min_cells_per_condition)
+    for candidate in candidate_conditions:
+        genes = parse_condition(candidate)
+        perturbation_type = "single" if len(genes) == 1 else "double"
+        gene_a = genes[0]
+        gene_b = genes[1] if len(genes) > 1 else ""
+
+        summary = candidate_summary[candidate]
+        n_cells = int(summary["n_cells"])
+        is_observed = int(summary["is_observed"] and n_cells >= min_cells)
+        if not is_observed:
+            n_cells = 0
+
+        rows.append(
+            {
+                "condition": candidate,
+                "perturbation_type": perturbation_type,
+                "gene_a": gene_a,
+                "gene_b": gene_b,
+                "perturb_gene": candidate if perturbation_type == "double" else gene_a,
+                "n_cells": n_cells,
+                "is_observed": is_observed,
+                "n_raw_conditions_merged": int(len(summary["raw_conditions"])) if is_observed else 0,
+                "raw_conditions_json": json.dumps(summary["raw_conditions"] if is_observed else [], ensure_ascii=False),
+                "condition_name_json": json.dumps(summary["condition_names"] if is_observed else [], ensure_ascii=False),
+            }
+        )
+
+    benchmark_df = pd.DataFrame(rows)
+    benchmark_df = benchmark_df.sort_values(["perturbation_type", "condition"]).reset_index(drop=True)
+    benchmark_df["class_id"] = np.arange(len(benchmark_df), dtype=np.int64)
+    benchmark_df["class_id"] = benchmark_df["class_id"].astype(int)
+
+    observed_df = benchmark_df[benchmark_df["is_observed"] == 1].copy()
+    observed_single = int((observed_df["perturbation_type"] == "single").sum())
+    observed_double = int((observed_df["perturbation_type"] == "double").sum())
+
+    summary = {
+        "n_cells_total": int(adata.n_obs),
+        "n_genes_total": int(adata.n_vars),
+        "protocol": args.protocol or None,
+        "selected_genes": selected_genes,
+        "n_selected_genes": int(len(selected_genes)),
+        "n_candidate_conditions": int(len(benchmark_df)),
+        "n_candidate_single": int((benchmark_df["perturbation_type"] == "single").sum()),
+        "n_candidate_double": int((benchmark_df["perturbation_type"] == "double").sum()),
+        "n_observed_conditions": int(len(observed_df)),
+        "n_observed_single": observed_single,
+        "n_observed_double": observed_double,
+        "min_cells_per_condition": min_cells,
+        "condition_col": args.condition_col,
+        "condition_name_col": args.condition_name_col,
+        "alignment_note": (
+            f"Built from protocol '{args.protocol}'." if args.protocol
+            else "Built from custom gene pool."
+        ),
+    }
+
+    if protocol_split is not None:
+        with (output_dir / "protocol_split.json").open("w", encoding="utf-8") as handle:
+            json.dump(protocol_split, handle, indent=2, ensure_ascii=False)
+
+        protocol_conditions = set()
+        for split_name in ("train", "val", "test"):
+            for cond in protocol_split.get(split_name, []):
+                protocol_conditions.add(canonical_condition_name(cond))
+        observed_protocol = protocol_conditions & set(
+            observed_df["condition"].astype(str).tolist()
+        )
+        missing_in_data = sorted(protocol_conditions - observed_protocol)
+        summary["n_protocol_conditions"] = int(len(protocol_conditions))
+        summary["n_protocol_conditions_observed"] = int(len(observed_protocol))
+        summary["protocol_conditions_missing_in_data"] = missing_in_data
+        if missing_in_data:
+            print_sys(
+                f"WARNING: {len(missing_in_data)} protocol conditions missing in Norman data: {missing_in_data}"
+            )
+
+    benchmark_df.to_csv(output_dir / "gene_inverse_conditions.csv", index=False)
+    benchmark_df[
+        ["class_id", "condition", "perturbation_type", "gene_a", "gene_b", "is_observed", "n_cells"]
+    ].to_csv(output_dir / "class_mapping.csv", index=False)
+    with (output_dir / "benchmark_summary.json").open("w", encoding="utf-8") as handle:
+        json.dump(summary, handle, indent=2, ensure_ascii=False)
+
+    print_sys(json.dumps(summary, indent=2, ensure_ascii=False))
+    print_sys(f"Saved benchmark files to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

BioMod/gene_inverse_norman/common.py

@@ -0,0 +1,660 @@
+import json
+import random
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.distributed as dist
+from scipy import sparse
+from torch import nn
+from torch.utils.data import Dataset
+
+CURRENT_DIR = Path(__file__).resolve().parent
+if str(CURRENT_DIR) not in sys.path:
+    sys.path.insert(0, str(CURRENT_DIR))
+
+from utils import canonical_condition_name, parse_condition, print_sys, reindex_tensor_universal
+
+TOPK_METRICS: Tuple[int, ...] = tuple(range(1, 11))
+
+
+def seed_all(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def parse_raw_conditions_json(raw_value: object, fallback_condition: str) -> List[str]:
+    if isinstance(raw_value, list):
+        return [str(item) for item in raw_value]
+
+    if pd.isna(raw_value):
+        return [fallback_condition]
+
+    raw_str = str(raw_value).strip()
+    if not raw_str:
+        return [fallback_condition]
+
+    try:
+        parsed = json.loads(raw_str)
+        if isinstance(parsed, list) and parsed:
+            return [str(item) for item in parsed]
+    except json.JSONDecodeError:
+        pass
+    return [fallback_condition]
+
+
+def normalize_benchmark_table(benchmark_df: pd.DataFrame, min_cells_per_condition: int = 1) -> pd.DataFrame:
+    if "condition" not in benchmark_df.columns:
+        raise KeyError("benchmark_csv must contain column 'condition'")
+
+    table = benchmark_df.copy()
+    table["condition"] = table["condition"].astype(str)
+    table = table.drop_duplicates(subset=["condition"]).reset_index(drop=True)
+
+    if "n_cells" not in table.columns:
+        table["n_cells"] = 0
+    table["n_cells"] = table["n_cells"].fillna(0).astype(int)
+
+    if "is_observed" not in table.columns:
+        table["is_observed"] = (table["n_cells"] > 0).astype(int)
+    table["is_observed"] = table["is_observed"].fillna(0).astype(int)
+
+    if "perturbation_type" not in table.columns:
+        table["perturbation_type"] = table["condition"].apply(
+            lambda condition: "single" if len(parse_condition(condition)) == 1 else "double"
+        )
+    table["perturbation_type"] = table["perturbation_type"].astype(str)
+
+    if "gene_a" not in table.columns:
+        table["gene_a"] = table["condition"].apply(lambda condition: parse_condition(condition)[0])
+    if "gene_b" not in table.columns:
+        table["gene_b"] = table["condition"].apply(
+            lambda condition: parse_condition(condition)[1] if len(parse_condition(condition)) > 1 else ""
+        )
+    table["gene_a"] = table["gene_a"].fillna("").astype(str)
+    table["gene_b"] = table["gene_b"].fillna("").astype(str)
+
+    if "perturb_gene" not in table.columns:
+        table["perturb_gene"] = table.apply(
+            lambda row: row["condition"] if row["perturbation_type"] == "double" else row["gene_a"],
+            axis=1,
+        )
+    table["perturb_gene"] = table["perturb_gene"].astype(str)
+
+    if min_cells_per_condition > 1:
+        observed_mask = (table["is_observed"] == 0) | (table["n_cells"] >= int(min_cells_per_condition))
+        table = table[observed_mask].copy()
+        table.loc[table["n_cells"] < int(min_cells_per_condition), "is_observed"] = 0
+        table.loc[table["n_cells"] < int(min_cells_per_condition), "n_cells"] = 0
+
+    if "class_id" in table.columns:
+        table = table.sort_values("class_id").reset_index(drop=True)
+    else:
+        table = table.sort_values(["perturbation_type", "condition"]).reset_index(drop=True)
+    table["class_id"] = np.arange(len(table), dtype=np.int64)
+    table["class_id"] = table["class_id"].astype(int)
+
+    return table
+
+
+def build_raw_to_canonical_map(benchmark_df: pd.DataFrame) -> Dict[str, str]:
+    raw_to_canonical: Dict[str, str] = {}
+    for _, row in benchmark_df.iterrows():
+        canonical = str(row["condition"])
+        raw_conditions = parse_raw_conditions_json(row.get("raw_conditions_json", np.nan), canonical)
+        for raw in raw_conditions:
+            raw_to_canonical[str(raw)] = canonical
+        raw_to_canonical[canonical] = canonical
+    return raw_to_canonical
+
+
+def collect_labeled_cells(
+    adata,
+    condition_col: str,
+    raw_to_canonical: Dict[str, str],
+    canonical_to_id: Dict[str, int],
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    if condition_col not in adata.obs.columns:
+        raise KeyError(f"condition_col '{condition_col}' not found in adata.obs")
+
+    condition_series = adata.obs[condition_col].astype(str)
+    canonical_series = condition_series.map(raw_to_canonical)
+    keep_mask = canonical_series.notna().to_numpy()
+    row_indices = np.where(keep_mask)[0]
+    canonical_labels = canonical_series.iloc[row_indices].astype(str).to_numpy()
+    class_ids = np.array([canonical_to_id[label] for label in canonical_labels.tolist()], dtype=np.int64)
+    return row_indices, class_ids, canonical_labels
+
+
+def split_observed_conditions(
+    class_df: pd.DataFrame,
+    seed: int,
+    train_cases: int,
+    val_cases: int,
+    test_cases: int,
+    allow_observed_count_mismatch: bool = False,
+) -> Dict[str, List[str]]:
+    observed_conditions = class_df.loc[class_df["is_observed"] == 1, "condition"].astype(str).tolist()
+    expected_cases = int(train_cases) + int(val_cases) + int(test_cases)
+    observed_count = len(observed_conditions)
+
+    if observed_count < expected_cases:
+        raise ValueError(
+            f"Observed perturbation cases ({observed_count}) are fewer than requested split size ({expected_cases})."
+        )
+    if observed_count != expected_cases and not allow_observed_count_mismatch:
+        raise ValueError(
+            "Observed perturbation case count does not match the requested CellFM-style split. "
+            f"observed={observed_count}, expected={expected_cases}. "
+            "Use the exact CellFM gene pool or pass --allow_observed_count_mismatch to continue."
+        )
+
+    rng = np.random.default_rng(seed)
+    observed_conditions = observed_conditions.copy()
+    rng.shuffle(observed_conditions)
+
+    usable_conditions = observed_conditions[:expected_cases]
+    unused_conditions = observed_conditions[expected_cases:]
+
+    train_conditions = sorted(usable_conditions[:train_cases])
+    val_conditions = sorted(usable_conditions[train_cases : train_cases + val_cases])
+    test_conditions = sorted(usable_conditions[train_cases + val_cases : expected_cases])
+
+    return {
+        "train": train_conditions,
+        "val": val_conditions,
+        "test": test_conditions,
+        "unused_observed": sorted(unused_conditions),
+    }
+
+
+def split_from_protocol(
+    class_df: pd.DataFrame,
+    protocol_split: Dict[str, object],
+    require_all_observed: bool = True,
+) -> Dict[str, List[str]]:
+    observed_set = set(class_df.loc[class_df["is_observed"] == 1, "condition"].astype(str).tolist())
+    candidate_set = set(class_df["condition"].astype(str).tolist())
+
+    def canonicalize_list(items: Sequence[str]) -> List[str]:
+        return [canonical_condition_name(str(item)) for item in items]
+
+    train = canonicalize_list(protocol_split.get("train", []))
+    val = canonicalize_list(protocol_split.get("val", []))
+    test = canonicalize_list(protocol_split.get("test", []))
+
+    all_listed = train + val + test
+    missing_in_candidates = sorted(set(all_listed) - candidate_set)
+    if missing_in_candidates:
+        raise ValueError(
+            f"Protocol conditions missing from benchmark candidate space: {missing_in_candidates}"
+        )
+
+    missing_in_data = sorted(set(all_listed) - observed_set)
+    if missing_in_data and require_all_observed:
+        raise ValueError(
+            "Protocol conditions are not observed in the local Norman data: "
+            f"{missing_in_data}. Re-run benchmark or pass --allow_protocol_unobserved."
+        )
+
+    return {
+        "train": sorted(train),
+        "val": sorted(val),
+        "test": sorted(test),
+        "unused_observed": sorted(observed_set - set(all_listed)),
+        "missing_in_data": missing_in_data,
+    }
+
+
+def extract_candidate_gene_names(class_df: pd.DataFrame) -> List[str]:
+    genes = set(class_df["gene_a"].astype(str).tolist())
+    genes.update(gene for gene in class_df["gene_b"].astype(str).tolist() if gene)
+    genes.discard("")
+    return sorted(genes)
+
+
+class GeneInverseCellDataset(Dataset):
+    def __init__(self, matrix, row_indices: np.ndarray, class_ids: np.ndarray):
+        if row_indices.shape[0] != class_ids.shape[0]:
+            raise ValueError("row_indices and class_ids must have the same length")
+        self.matrix = matrix
+        self.row_indices = row_indices.astype(np.int64, copy=False)
+        self.class_ids = class_ids.astype(np.int64, copy=False)
+
+    def __len__(self) -> int:
+        return int(self.row_indices.shape[0])
+
+    def __getitem__(self, idx: int):
+        row_idx = int(self.row_indices[idx])
+        row = self.matrix[row_idx]
+        if sparse.issparse(row):
+            row = row.toarray().ravel()
+        else:
+            row = np.asarray(row).reshape(-1)
+        x = torch.from_numpy(row.astype(np.float32, copy=False))
+        y = torch.tensor(int(self.class_ids[idx]), dtype=torch.long)
+        return x, y, row_idx
+
+
+class GeneInverseCompositionalRetrievalModel(nn.Module):
+    def __init__(
+        self,
+        backbone: nn.Module,
+        class_df: pd.DataFrame,
+        pool_method: str = "max",
+        head_hidden_dim: int = 0,
+        head_dropout: float = 0.1,
+        head_residual: bool = True,
+        temperature: float = 0.07,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.pool_method = pool_method
+        self.head_residual = bool(head_residual)
+        self.temperature = float(temperature)
+
+        embed_dim = int(backbone.base_dim)
+        hidden_dim = head_hidden_dim if head_hidden_dim > 0 else embed_dim * 4
+        self.cell_head = nn.Sequential(
+            nn.LayerNorm(embed_dim),
+            nn.Dropout(head_dropout),
+            nn.Linear(embed_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(hidden_dim, embed_dim),
+        )
+
+        candidate_gene_names = extract_candidate_gene_names(class_df)
+        gene_to_index = {gene: idx for idx, gene in enumerate(candidate_gene_names)}
+        candidate_index = np.full((len(class_df), 2), -1, dtype=np.int64)
+        candidate_mask = np.zeros((len(class_df), 2), dtype=np.float32)
+        for row_idx, (_, row) in enumerate(class_df.iterrows()):
+            genes = [gene for gene in [str(row["gene_a"]), str(row["gene_b"])] if gene and gene != "nan"]
+            for gene_pos, gene in enumerate(genes[:2]):
+                candidate_index[row_idx, gene_pos] = gene_to_index[gene]
+                candidate_mask[row_idx, gene_pos] = 1.0
+
+        self.candidate_gene_names = candidate_gene_names
+        self.candidate_gene_embeddings = nn.Embedding(len(candidate_gene_names), embed_dim)
+        self.register_buffer("candidate_gene_indices", torch.tensor(candidate_index, dtype=torch.long), persistent=False)
+        self.register_buffer("candidate_gene_mask", torch.tensor(candidate_mask, dtype=torch.float32), persistent=False)
+
+    def _pool_cell_embedding(self, outputs: Dict[str, torch.Tensor]) -> torch.Tensor:
+        module_encodings = outputs["module_encodings"]
+
+        if self.pool_method == "mean":
+            return module_encodings.mean(dim=1)
+        if self.pool_method == "max":
+            return module_encodings.max(dim=1)[0]
+        raise ValueError(f"Unsupported pool_method: {self.pool_method}. Expected one of ['mean', 'max'].")
+
+    def get_candidate_embeddings(self) -> torch.Tensor:
+        safe_indices = self.candidate_gene_indices.clamp_min(0)
+        gene_embeddings = self.candidate_gene_embeddings(safe_indices)
+        mask = self.candidate_gene_mask.unsqueeze(-1)
+        summed = (gene_embeddings * mask).sum(dim=1)
+        denom = mask.sum(dim=1).clamp_min(1.0)
+        candidate_embeddings = summed / denom
+        return nn.functional.normalize(candidate_embeddings, dim=-1)
+
+    def forward(self, x_model: torch.Tensor) -> torch.Tensor:
+        outputs = self.backbone(x_model, mask_ratio=0.0, return_encodings=True)
+        pooled_embedding = self._pool_cell_embedding(outputs)
+        cell_embedding = self.cell_head(pooled_embedding)
+        if self.head_residual:
+            cell_embedding = pooled_embedding + cell_embedding
+        cell_embedding = nn.functional.normalize(cell_embedding, dim=-1)
+        candidate_embeddings = self.get_candidate_embeddings()
+        logits = cell_embedding @ candidate_embeddings.T
+        return logits / self.temperature
+
+
+def set_requires_grad(module: nn.Module, flag: bool) -> None:
+    for param in module.parameters():
+        param.requires_grad = flag
+
+
+def configure_backbone_trainability(backbone: nn.Module, finetune_mode: str) -> Dict[str, int]:
+    set_requires_grad(backbone, False)
+
+    if finetune_mode == "full":
+        set_requires_grad(backbone.module_embed, True)
+        set_requires_grad(backbone.encoder, True)
+        set_requires_grad(backbone.encoder_norm, True)
+        backbone.pos_emb.requires_grad = True
+        if getattr(backbone, "cls_token", None) is not None:
+            backbone.cls_token.requires_grad = True
+    elif finetune_mode != "frozen":
+        raise ValueError(f"Unsupported finetune_mode: {finetune_mode}")
+
+    total_params = sum(param.numel() for param in backbone.parameters())
+    trainable_params = sum(param.numel() for param in backbone.parameters() if param.requires_grad)
+    return {"total": int(total_params), "trainable": int(trainable_params)}
+
+
+def build_optimizer(
+    model: nn.Module,
+    lr_head: float,
+    lr_backbone: float,
+    weight_decay: float,
+) -> torch.optim.Optimizer:
+    backbone_params = []
+    head_params = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if name.startswith("backbone."):
+            backbone_params.append(param)
+        else:
+            head_params.append(param)
+
+    param_groups = []
+    if backbone_params:
+        param_groups.append({"params": backbone_params, "lr": lr_backbone})
+    if head_params:
+        param_groups.append({"params": head_params, "lr": lr_head})
+    if not param_groups:
+        raise RuntimeError("No trainable parameters found")
+
+    return torch.optim.AdamW(
+        param_groups,
+        weight_decay=weight_decay,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+    )
+
+
+def compute_topk_metrics(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    topk_values: Sequence[int] = TOPK_METRICS,
+    prefix: str = "",
+) -> Dict[str, float]:
+    if logits.numel() == 0:
+        raise ValueError("logits is empty")
+
+    num_classes = int(logits.shape[1])
+    maxk = min(max(topk_values), num_classes)
+    topk_pred = logits.topk(maxk, dim=1).indices
+    targets_expanded = targets.unsqueeze(1)
+    correct = topk_pred.eq(targets_expanded)
+
+    metrics: Dict[str, float] = {}
+    metric_prefix = f"{prefix}_" if prefix else ""
+    for k in topk_values:
+        k_eff = min(k, num_classes)
+        value = correct[:, :k_eff].any(dim=1).float().mean().item()
+        metrics[f"{metric_prefix}top{k}_acc"] = float(value)
+    return metrics
+
+
+def compute_condition_hit_rates(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    condition_names: Sequence[str],
+    topk_values: Sequence[int] = TOPK_METRICS,
+) -> Tuple[Dict[str, float], torch.Tensor, torch.Tensor, List[str]]:
+    grouped_logits: Dict[str, List[torch.Tensor]] = {}
+    grouped_targets: Dict[str, List[int]] = {}
+    for logit_row, target, condition_name in zip(logits, targets.tolist(), condition_names):
+        grouped_logits.setdefault(str(condition_name), []).append(logit_row.detach().cpu())
+        grouped_targets.setdefault(str(condition_name), []).append(int(target))
+
+    aggregated_logits = []
+    aggregated_targets = []
+    ordered_conditions = sorted(grouped_logits.keys())
+    for condition_name in ordered_conditions:
+        target_ids = sorted(set(grouped_targets[condition_name]))
+        if len(target_ids) != 1:
+            raise ValueError(f"Condition {condition_name} maps to multiple targets: {target_ids}")
+        aggregated_targets.append(target_ids[0])
+        stacked = torch.stack(grouped_logits[condition_name], dim=0)
+        aggregated_logits.append(stacked.mean(dim=0))
+
+    aggregated_logits_tensor = torch.stack(aggregated_logits, dim=0)
+    aggregated_targets_tensor = torch.tensor(aggregated_targets, dtype=torch.long)
+    metrics = compute_topk_metrics(
+        aggregated_logits_tensor,
+        aggregated_targets_tensor,
+        topk_values=topk_values,
+        prefix="condition",
+    )
+    metrics["n_conditions"] = float(len(ordered_conditions))
+    return metrics, aggregated_logits_tensor, aggregated_targets_tensor, ordered_conditions
+
+
+def compute_condition_relevant_hit_rates(
+    aggregated_logits: torch.Tensor,
+    aggregated_targets: torch.Tensor,
+    id_to_condition: Sequence[str],
+    topk_values: Sequence[int] = TOPK_METRICS,
+    prefix: str = "condition_relevant",
+) -> Dict[str, float]:
+    if aggregated_logits.numel() == 0:
+        raise ValueError("aggregated_logits is empty")
+    if aggregated_logits.shape[0] != aggregated_targets.shape[0]:
+        raise ValueError(
+            "aggregated_logits and aggregated_targets must have matching first dimension"
+        )
+
+    num_classes = int(aggregated_logits.shape[1])
+    maxk = min(max(topk_values), num_classes)
+    topk_pred = aggregated_logits.topk(maxk, dim=1).indices.cpu().numpy()
+    target_ids = aggregated_targets.cpu().numpy()
+
+    relevant_hits = {int(k): 0 for k in topk_values}
+    for row_idx, true_id in enumerate(target_ids.tolist()):
+        true_genes = set(parse_condition(id_to_condition[int(true_id)]))
+        for k in topk_values:
+            k_eff = min(int(k), num_classes)
+            pred_ids = topk_pred[row_idx, :k_eff].tolist()
+            hit = any(
+                bool(true_genes & set(parse_condition(id_to_condition[int(pred_id)])))
+                for pred_id in pred_ids
+            )
+            if hit:
+                relevant_hits[int(k)] += 1
+
+    n_conditions = max(int(aggregated_logits.shape[0]), 1)
+    metrics: Dict[str, float] = {}
+    for k in topk_values:
+        metrics[f"{prefix}_top{int(k)}_acc"] = float(relevant_hits[int(k)] / n_conditions)
+    return metrics
+
+
+def compute_case_cellmean_hit_rates(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    condition_names: Sequence[str],
+    id_to_condition: Sequence[str],
+    topk_values: Sequence[int] = TOPK_METRICS,
+    prefix: str = "case_cellmean",
+) -> Dict[str, float]:
+    if logits.numel() == 0:
+        raise ValueError("logits is empty")
+    if logits.shape[0] != targets.shape[0]:
+        raise ValueError("logits and targets must have matching first dimension")
+    if logits.shape[0] != len(condition_names):
+        raise ValueError("condition_names length must match number of logits rows")
+
+    num_classes = int(logits.shape[1])
+    maxk = min(max(topk_values), num_classes)
+    topk_pred = logits.topk(maxk, dim=1).indices.cpu().numpy()
+    target_ids = targets.cpu().numpy().tolist()
+
+    condition_counts: Dict[str, int] = {}
+    condition_exact_hits: Dict[str, Dict[int, int]] = {}
+    condition_relevant_hits: Dict[str, Dict[int, int]] = {}
+
+    for row_idx, condition_name in enumerate(condition_names):
+        condition = str(condition_name)
+        condition_counts[condition] = condition_counts.get(condition, 0) + 1
+        if condition not in condition_exact_hits:
+            condition_exact_hits[condition] = {int(k): 0 for k in topk_values}
+            condition_relevant_hits[condition] = {int(k): 0 for k in topk_values}
+
+        true_id = int(target_ids[row_idx])
+        true_genes = set(parse_condition(id_to_condition[true_id]))
+        for k in topk_values:
+            k_eff = min(int(k), num_classes)
+            pred_ids = topk_pred[row_idx, :k_eff].tolist()
+            exact_hit = true_id in pred_ids
+            relevant_hit = any(
+                bool(true_genes & set(parse_condition(id_to_condition[int(pred_id)])))
+                for pred_id in pred_ids
+            )
+            if exact_hit:
+                condition_exact_hits[condition][int(k)] += 1
+            if relevant_hit:
+                condition_relevant_hits[condition][int(k)] += 1
+
+    ordered_conditions = sorted(condition_counts.keys())
+    n_conditions = max(len(ordered_conditions), 1)
+    metrics: Dict[str, float] = {}
+    for k in topk_values:
+        exact_case_rates = [
+            float(condition_exact_hits[condition][int(k)] / condition_counts[condition])
+            for condition in ordered_conditions
+        ]
+        relevant_case_rates = [
+            float(condition_relevant_hits[condition][int(k)] / condition_counts[condition])
+            for condition in ordered_conditions
+        ]
+        metrics[f"{prefix}_top{int(k)}_acc"] = float(np.mean(exact_case_rates))
+        metrics[f"{prefix}_relevant_top{int(k)}_acc"] = float(np.mean(relevant_case_rates))
+    metrics[f"{prefix}_n_conditions"] = float(n_conditions)
+    return metrics
+
+
+def build_condition_ranking_records(
+    aggregated_logits: torch.Tensor,
+    aggregated_targets: torch.Tensor,
+    ordered_conditions: Sequence[str],
+    id_to_condition: Sequence[str],
+    top_k: int = 10,
+) -> List[Dict[str, object]]:
+    records: List[Dict[str, object]] = []
+    top_indices = aggregated_logits.topk(min(top_k, aggregated_logits.shape[1]), dim=1).indices.cpu().numpy()
+    sorted_indices = aggregated_logits.argsort(dim=1, descending=True).cpu().numpy()
+
+    for idx, condition_name in enumerate(ordered_conditions):
+        true_id = int(aggregated_targets[idx].item())
+        ranking = sorted_indices[idx]
+        true_rank = int(np.where(ranking == true_id)[0][0]) + 1
+        pred_ids = top_indices[idx].tolist()
+        records.append(
+            {
+                "condition": str(condition_name),
+                "true_class_id": true_id,
+                "true_condition": str(id_to_condition[true_id]),
+                "true_rank": true_rank,
+                "top_predictions": [str(id_to_condition[pred_id]) for pred_id in pred_ids],
+            }
+        )
+    return records
+
+
+def evaluate_loader(
+    model: nn.Module,
+    dataloader,
+    device: torch.device,
+    scfm_index_positions: Sequence[int],
+    loss_fn: Optional[nn.Module] = None,
+    return_outputs: bool = False,
+    is_distributed: bool = False,
+    world_size: int = 1,
+    rank: int = 0,
+) -> Tuple[Dict[str, float], Optional[torch.Tensor], Optional[torch.Tensor], Optional[np.ndarray]]:
+    model.eval()
+    total_loss = 0.0
+    total_samples = 0
+    logits_all = []
+    targets_all = []
+    row_ids_all = []
+
+    with torch.no_grad():
+        for x, y, row_ids in dataloader:
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+            x_model = reindex_tensor_universal(x, scfm_index_positions, dim=1, filler="0", device=device)
+            logits = model(x_model)
+
+            if loss_fn is not None:
+                loss = loss_fn(logits, y)
+                total_loss += float(loss.item()) * int(x.shape[0])
+            total_samples += int(x.shape[0])
+
+            logits_all.append(logits.detach().cpu())
+            targets_all.append(y.detach().cpu())
+            row_ids_all.append(np.asarray(row_ids, dtype=np.int64))
+
+    if logits_all:
+        logits_cat = torch.cat(logits_all, dim=0)
+        targets_cat = torch.cat(targets_all, dim=0)
+        row_ids_cat = np.concatenate(row_ids_all, axis=0)
+    else:
+        logits_cat = torch.empty((0, 0), dtype=torch.float32)
+        targets_cat = torch.empty((0,), dtype=torch.long)
+        row_ids_cat = np.empty((0,), dtype=np.int64)
+
+    if is_distributed:
+        loss_tensor = torch.tensor([total_loss], device=device)
+        samples_tensor = torch.tensor([total_samples], device=device)
+        dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
+        dist.all_reduce(samples_tensor, op=dist.ReduceOp.SUM)
+        total_loss = float(loss_tensor.item())
+        total_samples = int(samples_tensor.item())
+
+        gather_payload = {
+            "logits": logits_cat.numpy(),
+            "targets": targets_cat.numpy(),
+            "row_ids": row_ids_cat,
+        }
+        gathered_payloads = [None for _ in range(world_size)]
+        dist.all_gather_object(gathered_payloads, gather_payload)
+
+        if rank != 0:
+            return {}, None, None, None
+
+        logits_list = [torch.from_numpy(payload["logits"]).float() for payload in gathered_payloads if payload["logits"].size > 0]
+        targets_list = [torch.from_numpy(payload["targets"]).long() for payload in gathered_payloads if payload["targets"].size > 0]
+        row_list = [payload["row_ids"].astype(np.int64) for payload in gathered_payloads if payload["row_ids"].size > 0]
+
+        if logits_list:
+            logits_cat = torch.cat(logits_list, dim=0)
+            targets_cat = torch.cat(targets_list, dim=0)
+            row_ids_cat = np.concatenate(row_list, axis=0)
+        else:
+            logits_cat = torch.empty((0, 0), dtype=torch.float32)
+            targets_cat = torch.empty((0,), dtype=torch.long)
+            row_ids_cat = np.empty((0,), dtype=np.int64)
+
+    metrics = compute_topk_metrics(logits_cat, targets_cat, topk_values=TOPK_METRICS)
+    metrics["loss"] = float(total_loss / max(total_samples, 1)) if loss_fn is not None else 0.0
+
+    if return_outputs:
+        return metrics, logits_cat, targets_cat, row_ids_cat
+    return metrics, None, None, None
+
+
+def fallback_gene_col(adata, gene_col: str) -> str:
+    if gene_col in adata.var.columns:
+        return gene_col
+    fallback_col = "_gene_inverse_var_index"
+    adata.var[fallback_col] = adata.var.index.astype(str)
+    print_sys(
+        f"gene_col '{gene_col}' not found in adata.var.columns; "
+        f"fallback to adata.var.index via temporary column '{fallback_col}'."
+    )
+    return fallback_col
+
+
+def rebuild_canonical_from_raw(raw_condition: str, collapse_ctrl_orientation: bool) -> str:
+    if collapse_ctrl_orientation:
+        return canonical_condition_name(raw_condition)
+    return str(raw_condition)

BioMod/gene_inverse_norman/evaluate_gene_inverse.py

@@ -0,0 +1,405 @@
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Dict, List
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+CURRENT_DIR = Path(__file__).resolve().parent
+if str(CURRENT_DIR) not in sys.path:
+    sys.path.insert(0, str(CURRENT_DIR))
+
+from common import (
+    GeneInverseCellDataset,
+    GeneInverseCompositionalRetrievalModel,
+    build_condition_ranking_records,
+    build_raw_to_canonical_map,
+    collect_labeled_cells,
+    compute_case_cellmean_hit_rates,
+    compute_condition_hit_rates,
+    compute_condition_relevant_hit_rates,
+    evaluate_loader,
+    fallback_gene_col,
+    normalize_benchmark_table,
+    seed_all,
+)
+from utils import (
+    default_repo_path,
+    load_pretrained_backbone,
+    print_sys,
+    resolve_input_gene_ids,
+)
+
+
+TOPK_RANGE = list(range(1, 11))
+
+
+def build_topk_table_rows(metrics: Dict[str, float], split_name: str, algorithm_name: str) -> pd.DataFrame:
+    families = [
+        ("top", "cell_correct"),
+        ("condition_top", "condition_correct"),
+        ("condition_relevant_top", "condition_relevant"),
+        ("case_cellmean_top", "case_cellmean_correct"),
+        ("case_cellmean_relevant_top", "case_cellmean_relevant"),
+    ]
+    rows = []
+    for key_prefix, family_name in families:
+        row = {
+            "algorithm": str(algorithm_name),
+            "split": str(split_name),
+            "metric_family": family_name,
+        }
+        has_value = False
+        for k in TOPK_RANGE:
+            key = f"{key_prefix}{k}_acc"
+            value = metrics.get(key)
+            if value is not None:
+                has_value = True
+            row[f"top{k}"] = float(value) if value is not None else np.nan
+        if has_value:
+            rows.append(row)
+    return pd.DataFrame(rows)
+
+
+def print_topk_table(metrics_df: pd.DataFrame, title: str) -> None:
+    if metrics_df.empty:
+        return
+    printable_df = metrics_df.copy()
+    for k in TOPK_RANGE:
+        col = f"top{k}"
+        printable_df[col] = printable_df[col].map(lambda x: "" if pd.isna(x) else f"{x:.4f}")
+    print_sys(title)
+    print_sys(printable_df.to_string(index=False))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Evaluate a CellFM-style reverse-perturbation retrieval head and optionally export logits."
+    )
+    parser.add_argument("--trained_head_ckpt", type=str, required=True, help="Path to best_finetune_gene_inverse.pth")
+    parser.add_argument("--scfm_hyper_params_path", type=str, required=True, help="Path to pretrained PKL")
+    parser.add_argument("--scfm_ckpt_path", type=str, required=True, help="Path to pretrained checkpoint")
+    parser.add_argument(
+        "--gene2vec_path",
+        type=str,
+        default=default_repo_path("selected_gene2vec_27k.npy"),
+        help="Path to gene2vec file",
+    )
+    parser.add_argument(
+        "--gene_mapping_path",
+        type=str,
+        default=default_repo_path("gene2ens.csv"),
+        help="Path to gene symbol -> Ensembl CSV",
+    )
+    parser.add_argument(
+        "--scfm_genes_file",
+        type=str,
+        default=default_repo_path("selected_genes_27k.txt"),
+        help="Path to selected model genes",
+    )
+    parser.add_argument(
+        "--adata_path",
+        type=str,
+        default=default_repo_path("GEARS", "data", "norman", "perturb_processed.h5ad"),
+        help="Path to Norman h5ad",
+    )
+    parser.add_argument(
+        "--benchmark_csv",
+        type=str,
+        default=default_repo_path(
+            "performer_module",
+            "gene_inverse_norman",
+            "outputs",
+            "norman_gene_inverse_benchmark",
+            "gene_inverse_conditions.csv",
+        ),
+        help="Benchmark CSV produced by build_norman_gene_inverse_benchmark.py",
+    )
+    parser.add_argument("--condition_col", type=str, default="condition")
+    parser.add_argument("--gene_col", type=str, default="gene_name")
+    parser.add_argument("--min_cells_per_condition", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--split",
+        type=str,
+        default="test",
+        choices=["train", "val", "test", "all"],
+        help="Which split to evaluate.",
+    )
+    parser.add_argument(
+        "--split_indices_npz",
+        type=str,
+        default="",
+        help="Optional split_indices.npz from finetune run for cell rows.",
+    )
+    parser.add_argument(
+        "--condition_splits_npz",
+        type=str,
+        default="",
+        help="Optional condition_splits.npz from finetune run for perturbation-level splits.",
+    )
+    parser.add_argument(
+        "--export_logits",
+        action="store_true",
+        help="Export per-cell raw logits CSV.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=default_repo_path("performer_module", "gene_inverse_norman", "outputs", "norman_gene_inverse_eval"),
+    )
+    parser.add_argument(
+        "--algorithm_name",
+        type=str,
+        default="",
+        help="Optional algorithm label in output tables; falls back to checkpoint value or scLong.",
+    )
+    args = parser.parse_args()
+
+    seed_all(args.seed)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    head_ckpt = torch.load(args.trained_head_ckpt, map_location="cpu")
+    run_config = head_ckpt.get("run_config", {})
+    pool_method = str(head_ckpt.get("pool_method", run_config.get("pool_method", "max")))
+    if pool_method not in {"mean", "max"}:
+        print_sys(f"pool_method '{pool_method}' is deprecated in current evaluator; fallback to 'max'.")
+        pool_method = "max"
+    head_hidden_dim = int(head_ckpt.get("head_hidden_dim", run_config.get("head_hidden_dim", 0)))
+    head_dropout = float(head_ckpt.get("head_dropout", run_config.get("head_dropout", 0.1)))
+    head_residual = bool(head_ckpt.get("head_residual", run_config.get("head_residual", False)))
+    temperature = float(head_ckpt.get("temperature", run_config.get("temperature", 0.07)))
+    algorithm_name = args.algorithm_name or str(run_config.get("algorithm_name", "scLong"))
+
+    class_table = head_ckpt.get("class_table")
+    if class_table:
+        class_df = normalize_benchmark_table(
+            benchmark_df=pd.DataFrame(class_table),
+            min_cells_per_condition=1,
+        )
+        print_sys("Using class_table from trained checkpoint for label mapping.")
+    else:
+        benchmark_df = pd.read_csv(args.benchmark_csv)
+        class_df = normalize_benchmark_table(
+            benchmark_df=benchmark_df,
+            min_cells_per_condition=args.min_cells_per_condition,
+        )
+        print_sys("Checkpoint has no class_table; fallback to benchmark_csv for label mapping.")
+    class_df = class_df.sort_values("class_id").reset_index(drop=True)
+    canonical_to_id: Dict[str, int] = dict(zip(class_df["condition"].tolist(), class_df["class_id"].tolist()))
+    id_to_condition: List[str] = class_df["condition"].astype(str).tolist()
+    id_to_gene: List[str] = class_df["perturb_gene"].astype(str).tolist()
+
+    print_sys(f"Loading Norman data from {args.adata_path}")
+    adata = ad.read_h5ad(args.adata_path)
+    raw_to_canonical = build_raw_to_canonical_map(class_df[class_df["is_observed"] == 1])
+    labeled_rows, labeled_class_ids, canonical_labels = collect_labeled_cells(
+        adata=adata,
+        condition_col=args.condition_col,
+        raw_to_canonical=raw_to_canonical,
+        canonical_to_id=canonical_to_id,
+    )
+    if labeled_rows.size == 0:
+        raise RuntimeError("No labeled cells found after mapping conditions to classes")
+
+    row_to_class = dict(zip(labeled_rows.tolist(), labeled_class_ids.tolist()))
+    row_to_condition = dict(zip(labeled_rows.tolist(), canonical_labels.tolist()))
+
+    split_npz = args.split_indices_npz
+    if not split_npz:
+        candidate = Path(args.trained_head_ckpt).resolve().parent / "split_indices.npz"
+        if candidate.exists():
+            split_npz = str(candidate)
+
+    condition_splits_npz = args.condition_splits_npz
+    if not condition_splits_npz:
+        candidate = Path(args.trained_head_ckpt).resolve().parent / "condition_splits.npz"
+        if candidate.exists():
+            condition_splits_npz = str(candidate)
+
+    if split_npz:
+        split_data = np.load(split_npz, allow_pickle=True)
+        train_rows = split_data["train_rows"].astype(np.int64)
+        val_rows = split_data["val_rows"].astype(np.int64)
+        test_rows = split_data["test_rows"].astype(np.int64)
+        print_sys(f"Loaded cell split indices from {split_npz}")
+    elif condition_splits_npz:
+        condition_data = np.load(condition_splits_npz, allow_pickle=True)
+        train_conditions = set(condition_data["train_conditions"].tolist())
+        val_conditions = set(condition_data["val_conditions"].tolist())
+        test_conditions = set(condition_data["test_conditions"].tolist())
+        train_rows = np.array([row for row in labeled_rows.tolist() if row_to_condition[row] in train_conditions], dtype=np.int64)
+        val_rows = np.array([row for row in labeled_rows.tolist() if row_to_condition[row] in val_conditions], dtype=np.int64)
+        test_rows = np.array([row for row in labeled_rows.tolist() if row_to_condition[row] in test_conditions], dtype=np.int64)
+        print_sys(f"Loaded perturbation splits from {condition_splits_npz}")
+    else:
+        raise RuntimeError("No split artifacts found. Provide --split_indices_npz or --condition_splits_npz.")
+
+    if args.split == "train":
+        eval_rows = train_rows
+    elif args.split == "val":
+        eval_rows = val_rows
+    elif args.split == "test":
+        eval_rows = test_rows
+    else:
+        eval_rows = np.concatenate([train_rows, val_rows, test_rows], axis=0)
+
+    eval_labels = np.array([row_to_class[int(idx)] for idx in eval_rows], dtype=np.int64)
+    eval_condition_names = [row_to_condition[int(idx)] for idx in eval_rows]
+    eval_dataset = GeneInverseCellDataset(adata.X, eval_rows, eval_labels)
+    eval_loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    backbone, _ = load_pretrained_backbone(
+        scfm_hyper_params_path=args.scfm_hyper_params_path,
+        scfm_ckpt_path=args.scfm_ckpt_path,
+        gene2vec_path=args.gene2vec_path,
+        device=device,
+    )
+    model = GeneInverseCompositionalRetrievalModel(
+        backbone=backbone,
+        class_df=class_df,
+        pool_method=pool_method,
+        head_hidden_dim=head_hidden_dim,
+        head_dropout=head_dropout,
+        head_residual=head_residual,
+        temperature=temperature,
+    ).to(device)
+    model.load_state_dict(head_ckpt["model_state_dict"], strict=True)
+
+    gene2ens = pd.read_csv(args.gene_mapping_path, index_col=0)
+    symbol_to_ensembl = gene2ens["Ensembl gene ID"].to_dict()
+    with open(args.scfm_genes_file, "r", encoding="utf-8") as handle:
+        scfm_genes = [line.rstrip("\n") for line in handle.readlines()]
+    scfm_gene_set = set(scfm_genes)
+
+    resolved_gene_col = fallback_gene_col(adata, args.gene_col)
+    input_gene_ids = resolve_input_gene_ids(
+        adata=adata,
+        gene_col=resolved_gene_col,
+        symbol_to_ensembl=symbol_to_ensembl,
+        target_gene_set=scfm_gene_set,
+    )
+    input_mapping = {gene_id: index for index, gene_id in enumerate(input_gene_ids)}
+    input_gene_num = len(input_gene_ids)
+    scfm_genes_pad = scfm_genes + ["PAD"]
+    scfm_index_positions = [input_mapping.get(gene_id, input_gene_num) for gene_id in scfm_genes_pad]
+
+    loss_fn = nn.CrossEntropyLoss()
+    metrics, logits, targets, row_ids = evaluate_loader(
+        model=model,
+        dataloader=eval_loader,
+        device=device,
+        scfm_index_positions=scfm_index_positions,
+        loss_fn=loss_fn,
+        return_outputs=True,
+    )
+
+    pred_ids = logits.argmax(dim=1).numpy()
+    true_ids = targets.numpy()
+    pred_df = pd.DataFrame(
+        {
+            "row_index": row_ids.astype(int),
+            "cell_name": adata.obs_names[row_ids].astype(str).tolist(),
+            "true_class_id": true_ids.astype(int),
+            "pred_class_id": pred_ids.astype(int),
+            "true_condition": [id_to_condition[idx] for idx in true_ids],
+            "pred_condition": [id_to_condition[idx] for idx in pred_ids],
+            "true_gene": [id_to_gene[idx] for idx in true_ids],
+            "pred_gene": [id_to_gene[idx] for idx in pred_ids],
+            "correct": (pred_ids == true_ids).astype(int),
+        }
+    )
+    pred_df.to_csv(output_dir / f"{args.split}_predictions.csv", index=False)
+
+    if args.export_logits:
+        logits_df = pd.DataFrame(logits.numpy(), columns=id_to_condition)
+        logits_df.insert(0, "cell_name", adata.obs_names[row_ids].astype(str).tolist())
+        logits_df.insert(1, "row_index", row_ids.astype(int))
+        logits_df.to_csv(output_dir / f"{args.split}_logits.csv", index=False)
+
+    eval_condition_names_in_order = [row_to_condition[int(row)] for row in row_ids.tolist()]
+    condition_metrics, aggregated_logits, aggregated_targets, ordered_conditions = compute_condition_hit_rates(
+        logits=logits,
+        targets=targets,
+        condition_names=eval_condition_names_in_order,
+    )
+    metrics.update(condition_metrics)
+    relevant_metrics = compute_condition_relevant_hit_rates(
+        aggregated_logits=aggregated_logits,
+        aggregated_targets=aggregated_targets,
+        id_to_condition=id_to_condition,
+    )
+    metrics.update(relevant_metrics)
+    case_cellmean_metrics = compute_case_cellmean_hit_rates(
+        logits=logits,
+        targets=targets,
+        condition_names=eval_condition_names_in_order,
+        id_to_condition=id_to_condition,
+    )
+    metrics.update(case_cellmean_metrics)
+
+    ranking_records = build_condition_ranking_records(
+        aggregated_logits=aggregated_logits,
+        aggregated_targets=aggregated_targets,
+        ordered_conditions=ordered_conditions,
+        id_to_condition=id_to_condition,
+        top_k=10,
+    )
+    with (output_dir / f"{args.split}_condition_rankings.jsonl").open("w", encoding="utf-8") as handle:
+        for record in ranking_records:
+            handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+    summary = {
+        "algorithm_name": algorithm_name,
+        "protocol": run_config.get("protocol", None),
+        "split": args.split,
+        "n_cells": int(eval_rows.size),
+        "n_conditions": int(len(ordered_conditions)),
+        "n_classes": int(len(class_df)),
+        "pool_method": pool_method,
+        "head_hidden_dim": head_hidden_dim,
+        "head_dropout": head_dropout,
+        "head_residual": head_residual,
+        "temperature": temperature,
+        **{f"{args.split}_{k}": float(v) for k, v in metrics.items()},
+        "alignment_notes": [
+            "Primary CellFM-style comparison should use condition_top1_acc through condition_top10_acc.",
+            "Metrics aggregate cell logits to perturbation-level mean logits before computing condition hit rates.",
+            "scGPT-style relevant-hit metrics are also exported as condition_relevant_topk_acc.",
+            "Paper-style non-integer hit rates can be approximated via case_cellmean_topk_acc metrics.",
+            "Exact CellFM gene-pool identity must be verified separately if strict paper reproduction is required.",
+        ],
+    }
+    metrics_table_df = build_topk_table_rows(
+        metrics=metrics,
+        split_name=args.split,
+        algorithm_name=algorithm_name,
+    )
+    metrics_table_df.to_csv(output_dir / f"{args.split}_metrics_table.csv", index=False)
+    print_topk_table(metrics_table_df, title=f"Top-k Metrics Table ({args.split})")
+    with (output_dir / f"{args.split}_summary_metrics.json").open("w", encoding="utf-8") as handle:
+        json.dump(summary, handle, indent=2, ensure_ascii=False)
+
+    print_sys(json.dumps(summary, indent=2, ensure_ascii=False))
+    print_sys(f"Saved evaluation outputs to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

BioMod/gene_inverse_norman/finetune_gene_inverse.py

@@ -0,0 +1,837 @@
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Dict, List
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler
+
+CURRENT_DIR = Path(__file__).resolve().parent
+if str(CURRENT_DIR) not in sys.path:
+    sys.path.insert(0, str(CURRENT_DIR))
+
+from common import (
+    GeneInverseCellDataset,
+    GeneInverseCompositionalRetrievalModel,
+    build_condition_ranking_records,
+    build_optimizer,
+    build_raw_to_canonical_map,
+    collect_labeled_cells,
+    compute_case_cellmean_hit_rates,
+    compute_condition_hit_rates,
+    compute_condition_relevant_hit_rates,
+    configure_backbone_trainability,
+    evaluate_loader,
+    fallback_gene_col,
+    normalize_benchmark_table,
+    seed_all,
+    split_from_protocol,
+    split_observed_conditions,
+)
+from utils import (
+    default_repo_path,
+    load_pretrained_backbone,
+    print_sys,
+    reindex_tensor_universal,
+    resolve_input_gene_ids,
+)
+
+
+TOPK_RANGE = list(range(1, 11))
+
+
+def build_topk_table_rows(metrics: Dict[str, float], split_name: str, algorithm_name: str) -> pd.DataFrame:
+    families = [
+        ("top", "cell_correct"),
+        ("condition_top", "condition_correct"),
+        ("condition_relevant_top", "condition_relevant"),
+        ("case_cellmean_top", "case_cellmean_correct"),
+        ("case_cellmean_relevant_top", "case_cellmean_relevant"),
+    ]
+    rows = []
+    for key_prefix, family_name in families:
+        row = {
+            "algorithm": str(algorithm_name),
+            "split": str(split_name),
+            "metric_family": family_name,
+        }
+        has_value = False
+        for k in TOPK_RANGE:
+            key = f"{key_prefix}{k}_acc"
+            value = metrics.get(key)
+            if value is not None:
+                has_value = True
+            row[f"top{k}"] = float(value) if value is not None else np.nan
+        if has_value:
+            rows.append(row)
+    return pd.DataFrame(rows)
+
+
+def print_topk_table(metrics_df: pd.DataFrame, title: str) -> None:
+    if metrics_df.empty:
+        return
+    printable_df = metrics_df.copy()
+    for k in TOPK_RANGE:
+        col = f"top{k}"
+        printable_df[col] = printable_df[col].map(lambda x: "" if pd.isna(x) else f"{x:.4f}")
+    print_sys(title)
+    print_sys(printable_df.to_string(index=False))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Fine-tune performer_module for CellFM-style Norman reverse perturbation retrieval."
+    )
+    parser.add_argument("--scfm_hyper_params_path", type=str, required=True, help="Path to pretrained PKL")
+    parser.add_argument("--scfm_ckpt_path", type=str, required=True, help="Path to pretrained checkpoint")
+    parser.add_argument(
+        "--gene2vec_path",
+        type=str,
+        default=default_repo_path("selected_gene2vec_27k.npy"),
+        help="Path to gene2vec file",
+    )
+    parser.add_argument(
+        "--gene_mapping_path",
+        type=str,
+        default=default_repo_path("gene2ens.csv"),
+        help="Path to gene symbol -> Ensembl CSV",
+    )
+    parser.add_argument(
+        "--scfm_genes_file",
+        type=str,
+        default=default_repo_path("selected_genes_27k.txt"),
+        help="Path to selected model genes",
+    )
+    parser.add_argument(
+        "--adata_path",
+        type=str,
+        default=default_repo_path("GEARS", "data", "norman", "perturb_processed.h5ad"),
+        help="Path to Norman h5ad",
+    )
+    parser.add_argument(
+        "--benchmark_csv",
+        type=str,
+        default=default_repo_path(
+            "performer_module",
+            "gene_inverse_norman",
+            "outputs",
+            "norman_gene_inverse_benchmark",
+            "gene_inverse_conditions.csv",
+        ),
+        help="Benchmark CSV produced by build_norman_gene_inverse_benchmark.py",
+    )
+    parser.add_argument("--condition_col", type=str, default="condition", help="Condition column in adata.obs")
+    parser.add_argument("--gene_col", type=str, default="gene_name", help="Gene column in adata.var")
+    parser.add_argument("--min_cells_per_condition", type=int, default=1)
+    parser.add_argument("--epochs", type=int, default=5)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--num_workers", type=int, default=6)
+    parser.add_argument("--lr_head", type=float, default=1e-4)
+    parser.add_argument("--lr_backbone", type=float, default=1e-4)
+    parser.add_argument("--weight_decay", type=float, default=1e-3)
+    parser.add_argument(
+        "--class_balance_power",
+        type=float,
+        default=0.0,
+        help="If >0, use train class weights = count^(-power) in CE loss (e.g., 0.5 for inverse-sqrt).",
+    )
+    parser.add_argument(
+        "--label_smoothing",
+        type=float,
+        default=0.0,
+        help="Label smoothing for training CE loss. 0 disables it.",
+    )
+    parser.add_argument("--head_hidden_dim", type=int, default=0, help="0 means 4x embed dim")
+    parser.add_argument(
+        "--no_head_residual",
+        dest="head_residual",
+        action="store_false",
+        help="Disable residual connection in classification head.",
+    )
+    parser.add_argument("--head_dropout", type=float, default=0.1)
+    parser.add_argument("--temperature", type=float, default=0.07)
+    parser.add_argument(
+        "--pool_method",
+        type=str,
+        default="max",
+        choices=["mean", "max"],
+        help="Cell embedding pooling over module encodings.",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--finetune_mode", type=str, default="frozen", choices=["frozen", "full"])
+    parser.add_argument(
+        "--protocol",
+        type=str,
+        default="",
+        choices=["", "scgpt", "cellfm"],
+        help="Use a built-in protocol split (scgpt or cellfm). Overrides --train_cases/--val_cases/--test_cases.",
+    )
+    parser.add_argument("--train_cases", type=int, default=33)
+    parser.add_argument("--val_cases", type=int, default=3)
+    parser.add_argument("--test_cases", type=int, default=11)
+    parser.add_argument(
+        "--allow_observed_count_mismatch",
+        action="store_true",
+        help="Allow observed perturbation count to differ from train+val+test total when not using a protocol.",
+    )
+    parser.add_argument(
+        "--allow_protocol_unobserved",
+        action="store_true",
+        help="Allow protocol conditions to be missing in the local data (uses fewer cells).",
+    )
+    parser.add_argument(
+        "--local_rank",
+        type=int,
+        default=-1,
+        help="Local rank for distributed training. Set by torchrun automatically.",
+    )
+    parser.add_argument(
+        "--select_metric",
+        type=str,
+        default="condition_top10_acc",
+        choices=[f"condition_top{k}_acc" for k in range(1, 11)]
+        + [f"condition_relevant_top{k}_acc" for k in range(1, 11)]
+        + [f"case_cellmean_top{k}_acc" for k in range(1, 11)]
+        + [f"case_cellmean_relevant_top{k}_acc" for k in range(1, 11)]
+        + [f"top{k}_acc" for k in range(1, 11)],
+        help="Single validation metric used to select best checkpoint; all metrics are still reported.",
+    )
+    parser.add_argument("--early_stop_patience", type=int, default=0, help="0 disables early stopping")
+    parser.add_argument(
+        "--algorithm_name",
+        type=str,
+        default="scLong",
+        help="Name tag written to outputs for easier cross-algorithm comparison.",
+    )
+    parser.add_argument(
+        "--save_best_checkpoint",
+        action="store_true",
+        help="Save best_finetune_gene_inverse.pth. Disabled by default to reduce disk usage.",
+    )
+    parser.add_argument(
+        "--amp",
+        dest="amp",
+        action="store_true",
+        help="Enable AMP mixed precision (default).",
+    )
+    parser.add_argument(
+        "--no_amp",
+        dest="amp",
+        action="store_false",
+        help="Disable AMP mixed precision.",
+    )
+    parser.set_defaults(amp=True, head_residual=True)
+    parser.add_argument(
+        "--export_test_logits",
+        action="store_true",
+        help="Export per-cell raw logits CSV for the test split.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=default_repo_path("performer_module", "gene_inverse_norman", "outputs", "norman_gene_inverse_finetune"),
+        help="Output directory",
+    )
+    args = parser.parse_args()
+
+    is_distributed = args.local_rank != -1 or "LOCAL_RANK" in os.environ or "RANK" in os.environ
+    if is_distributed:
+        if "LOCAL_RANK" in os.environ:
+            local_rank = int(os.environ["LOCAL_RANK"])
+        elif "RANK" in os.environ:
+            local_rank = int(os.environ["RANK"]) % torch.cuda.device_count()
+        else:
+            local_rank = args.local_rank
+
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+        is_main_process = rank == 0
+        device = torch.device(f"cuda:{local_rank}")
+        torch.cuda.set_device(device)
+        print_sys(f"[DDP] Initialized: rank={rank}, local_rank={local_rank}, world_size={world_size}")
+    else:
+        world_size = 1
+        rank = 0
+        local_rank = 0
+        is_main_process = True
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print_sys(f"[Single GPU] Running on device: {device}")
+
+    seed_all(args.seed + rank)
+    output_dir = Path(args.output_dir)
+    if is_main_process:
+        output_dir.mkdir(parents=True, exist_ok=True)
+    if is_distributed:
+        dist.barrier()
+
+    benchmark_df = pd.read_csv(args.benchmark_csv)
+    class_df = normalize_benchmark_table(
+        benchmark_df=benchmark_df,
+        min_cells_per_condition=args.min_cells_per_condition,
+    )
+    if class_df.empty:
+        raise RuntimeError("No classes remain after benchmark filtering")
+
+    class_df = class_df.sort_values("class_id").reset_index(drop=True)
+    if is_main_process:
+        class_df.to_csv(output_dir / "class_mapping.csv", index=False)
+
+    canonical_to_id: Dict[str, int] = dict(zip(class_df["condition"].tolist(), class_df["class_id"].tolist()))
+    id_to_condition: List[str] = class_df.sort_values("class_id")["condition"].astype(str).tolist()
+    id_to_gene: List[str] = class_df.sort_values("class_id")["perturb_gene"].astype(str).tolist()
+
+    if is_main_process:
+        if args.protocol:
+            protocol_split_path = CURRENT_DIR / "protocols" / args.protocol / "split.json"
+            with open(protocol_split_path, "r", encoding="utf-8") as handle:
+                protocol_split = json.load(handle)
+            split_map = split_from_protocol(
+                class_df=class_df,
+                protocol_split=protocol_split,
+                require_all_observed=not args.allow_protocol_unobserved,
+            )
+            print_sys(
+                f"Using built-in protocol '{args.protocol}': "
+                f"train={len(split_map['train'])}, val={len(split_map['val'])}, test={len(split_map['test'])}"
+            )
+            if split_map.get("missing_in_data"):
+                print_sys(
+                    f"WARNING: protocol conditions missing in data: {split_map['missing_in_data']}"
+                )
+        else:
+            split_map = split_observed_conditions(
+                class_df=class_df,
+                seed=args.seed,
+                train_cases=args.train_cases,
+                val_cases=args.val_cases,
+                test_cases=args.test_cases,
+                allow_observed_count_mismatch=args.allow_observed_count_mismatch,
+            )
+        split_payload = {
+            "train_conditions": np.array(split_map["train"], dtype=object),
+            "val_conditions": np.array(split_map["val"], dtype=object),
+            "test_conditions": np.array(split_map["test"], dtype=object),
+            "unused_observed_conditions": np.array(split_map.get("unused_observed", []), dtype=object),
+        }
+        np.savez(output_dir / "condition_splits.npz", **split_payload)
+        with (output_dir / "condition_splits.json").open("w", encoding="utf-8") as handle:
+            json.dump(
+                {k: list(v) if isinstance(v, list) else v for k, v in split_map.items()},
+                handle,
+                indent=2,
+                ensure_ascii=False,
+            )
+    if is_distributed:
+        dist.barrier()
+
+    split_npz = np.load(output_dir / "condition_splits.npz", allow_pickle=True)
+    train_conditions = set(split_npz["train_conditions"].tolist())
+    val_conditions = set(split_npz["val_conditions"].tolist())
+    test_conditions = set(split_npz["test_conditions"].tolist())
+
+    print_sys(f"Loading Norman data from {args.adata_path}")
+    adata = ad.read_h5ad(args.adata_path)
+    raw_to_canonical = build_raw_to_canonical_map(class_df[class_df["is_observed"] == 1])
+    row_indices, class_ids, canonical_labels = collect_labeled_cells(
+        adata=adata,
+        condition_col=args.condition_col,
+        raw_to_canonical=raw_to_canonical,
+        canonical_to_id=canonical_to_id,
+    )
+    if row_indices.size == 0:
+        raise RuntimeError("No labeled cells found after mapping conditions to classes")
+
+    train_mask = np.array([label in train_conditions for label in canonical_labels], dtype=bool)
+    val_mask = np.array([label in val_conditions for label in canonical_labels], dtype=bool)
+    test_mask = np.array([label in test_conditions for label in canonical_labels], dtype=bool)
+
+    train_rows = row_indices[train_mask]
+    val_rows = row_indices[val_mask]
+    test_rows = row_indices[test_mask]
+    train_labels = class_ids[train_mask]
+    val_labels = class_ids[val_mask]
+    test_labels = class_ids[test_mask]
+
+    if is_main_process:
+        np.savez(
+            output_dir / "split_indices.npz",
+            train_rows=train_rows,
+            val_rows=val_rows,
+            test_rows=test_rows,
+        )
+
+    split_summary = {
+        "n_classes_total": int(len(class_df)),
+        "n_observed_classes": int((class_df["is_observed"] == 1).sum()),
+        "n_train_conditions": int(len(train_conditions)),
+        "n_val_conditions": int(len(val_conditions)),
+        "n_test_conditions": int(len(test_conditions)),
+        "n_train_cells": int(train_rows.size),
+        "n_val_cells": int(val_rows.size),
+        "n_test_cells": int(test_rows.size),
+    }
+    if is_main_process:
+        print_sys(json.dumps(split_summary, indent=2, ensure_ascii=False))
+
+    train_dataset = GeneInverseCellDataset(adata.X, train_rows, train_labels)
+    val_dataset = GeneInverseCellDataset(adata.X, val_rows, val_labels)
+    test_dataset = GeneInverseCellDataset(adata.X, test_rows, test_labels)
+
+    if is_distributed:
+        train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
+        val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank, shuffle=False)
+        test_sampler = DistributedSampler(test_dataset, num_replicas=world_size, rank=rank, shuffle=False)
+        train_shuffle = False
+    else:
+        train_sampler = None
+        val_sampler = None
+        test_sampler = None
+        train_shuffle = True
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=train_shuffle,
+        sampler=train_sampler,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=val_sampler,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=test_sampler,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    backbone, _ = load_pretrained_backbone(
+        scfm_hyper_params_path=args.scfm_hyper_params_path,
+        scfm_ckpt_path=args.scfm_ckpt_path,
+        gene2vec_path=args.gene2vec_path,
+        device=device,
+    )
+    trainability = configure_backbone_trainability(backbone=backbone, finetune_mode=args.finetune_mode)
+    model = GeneInverseCompositionalRetrievalModel(
+        backbone=backbone,
+        class_df=class_df,
+        pool_method=args.pool_method,
+        head_hidden_dim=args.head_hidden_dim,
+        head_dropout=args.head_dropout,
+        head_residual=args.head_residual,
+        temperature=args.temperature,
+    ).to(device)
+
+    if is_distributed:
+        model = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False)
+        unwrapped_model = model.module
+    else:
+        unwrapped_model = model
+
+    if is_main_process:
+        print_sys(
+            f"Backbone params: trainable={trainability['trainable']} total={trainability['total']} | "
+            f"Head params: trainable={sum(p.numel() for p in unwrapped_model.cell_head.parameters()) + sum(p.numel() for p in unwrapped_model.candidate_gene_embeddings.parameters())}"
+        )
+
+    gene2ens = pd.read_csv(args.gene_mapping_path, index_col=0)
+    symbol_to_ensembl = gene2ens["Ensembl gene ID"].to_dict()
+    with open(args.scfm_genes_file, "r", encoding="utf-8") as handle:
+        scfm_genes = [line.rstrip("\n") for line in handle.readlines()]
+    scfm_gene_set = set(scfm_genes)
+
+    resolved_gene_col = fallback_gene_col(adata, args.gene_col)
+    input_gene_ids = resolve_input_gene_ids(
+        adata=adata,
+        gene_col=resolved_gene_col,
+        symbol_to_ensembl=symbol_to_ensembl,
+        target_gene_set=scfm_gene_set,
+    )
+    input_mapping = {gene_id: index for index, gene_id in enumerate(input_gene_ids)}
+    input_gene_num = len(input_gene_ids)
+    scfm_genes_pad = scfm_genes + ["PAD"]
+    scfm_index_positions = [input_mapping.get(gene_id, input_gene_num) for gene_id in scfm_genes_pad]
+
+    optimizer = build_optimizer(
+        model=model,
+        lr_head=args.lr_head,
+        lr_backbone=args.lr_backbone,
+        weight_decay=args.weight_decay,
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(args.epochs, 1))
+    train_class_weights = None
+    if args.class_balance_power > 0:
+        train_counts = np.bincount(train_labels, minlength=len(class_df)).astype(np.float32)
+        class_weights = np.ones(len(class_df), dtype=np.float32)
+        nonzero = train_counts > 0
+        class_weights[nonzero] = np.power(train_counts[nonzero], -float(args.class_balance_power))
+        class_weights = class_weights / max(float(class_weights[nonzero].mean()), 1e-8)
+        train_class_weights = torch.tensor(class_weights, dtype=torch.float32, device=device)
+        if is_main_process:
+            print_sys(
+                f"Using class-balanced CE loss with power={args.class_balance_power:.3f} | "
+                f"weight range [{float(class_weights[nonzero].min()):.4f}, {float(class_weights[nonzero].max()):.4f}]"
+            )
+
+    train_loss_fn = nn.CrossEntropyLoss(
+        weight=train_class_weights,
+        label_smoothing=float(args.label_smoothing),
+    )
+    eval_loss_fn = nn.CrossEntropyLoss()
+
+    use_amp = bool(args.amp and device.type == "cuda")
+    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
+
+    history = []
+    best_score = -float("inf")
+    best_rank_tuple = (-float("inf"), -float("inf"), -float("inf"), -float("inf"))
+    best_epoch = -1
+    best_state = None
+    best_val_metrics = None
+    no_improve = 0
+
+    val_row_to_condition = {int(row): str(cond) for row, cond in zip(val_rows.tolist(), canonical_labels[val_mask].tolist())}
+    test_row_to_condition = {int(row): str(cond) for row, cond in zip(test_rows.tolist(), canonical_labels[test_mask].tolist())}
+
+    for epoch in range(1, args.epochs + 1):
+        if is_distributed and train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+
+        model.train()
+        running_loss = 0.0
+        running_samples = 0
+
+        for x, y, _ in train_loader:
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+            x_model = reindex_tensor_universal(x, scfm_index_positions, dim=1, filler="0", device=device)
+
+            optimizer.zero_grad(set_to_none=True)
+            with torch.amp.autocast("cuda", enabled=use_amp):
+                logits = model(x_model)
+                loss = train_loss_fn(logits, y)
+
+            if use_amp:
+                scaler.scale(loss).backward()
+                if args.grad_clip > 0:
+                    scaler.unscale_(optimizer)
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                loss.backward()
+                if args.grad_clip > 0:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+                optimizer.step()
+
+            batch_size = int(x.shape[0])
+            running_loss += float(loss.item()) * batch_size
+            running_samples += batch_size
+
+        train_loss = running_loss / max(running_samples, 1)
+        scheduler.step()
+        if is_distributed:
+            loss_tensor = torch.tensor([train_loss], device=device)
+            dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
+            train_loss = loss_tensor.item() / world_size
+
+        val_metrics, val_logits, val_targets, val_row_ids = evaluate_loader(
+            model=unwrapped_model,
+            dataloader=val_loader,
+            device=device,
+            scfm_index_positions=scfm_index_positions,
+            loss_fn=eval_loss_fn,
+            return_outputs=True,
+            is_distributed=is_distributed,
+            world_size=world_size,
+            rank=rank,
+        )
+
+        if is_main_process:
+            val_condition_names = [val_row_to_condition[int(row)] for row in val_row_ids.tolist()]
+            (
+                val_condition_metrics,
+                val_aggregated_logits,
+                val_aggregated_targets,
+                _,
+            ) = compute_condition_hit_rates(
+                logits=val_logits,
+                targets=val_targets,
+                condition_names=val_condition_names,
+            )
+            val_metrics.update(val_condition_metrics)
+            val_relevant_metrics = compute_condition_relevant_hit_rates(
+                aggregated_logits=val_aggregated_logits,
+                aggregated_targets=val_aggregated_targets,
+                id_to_condition=id_to_condition,
+            )
+            val_metrics.update(val_relevant_metrics)
+            val_case_cellmean_metrics = compute_case_cellmean_hit_rates(
+                logits=val_logits,
+                targets=val_targets,
+                condition_names=val_condition_names,
+                id_to_condition=id_to_condition,
+            )
+            val_metrics.update(val_case_cellmean_metrics)
+            record = {
+                "epoch": int(epoch),
+                "train_loss": float(train_loss),
+                **{f"val_{k}": float(v) for k, v in val_metrics.items()},
+            }
+            history.append(record)
+            print_sys(
+                f"Epoch {epoch} | train loss {train_loss:.4f} | "
+                f"val cond top1 {val_metrics['condition_top1_acc']:.4f} "
+                f"top5 {val_metrics['condition_top5_acc']:.4f} "
+                f"top10 {val_metrics['condition_top10_acc']:.4f} | "
+                f"val case top1 {val_metrics['case_cellmean_top1_acc']:.4f} "
+                f"top5 {val_metrics['case_cellmean_top5_acc']:.4f} | "
+                f"val rel top1 {val_metrics['condition_relevant_top1_acc']:.4f} | "
+                f"val cell top1 {val_metrics['top1_acc']:.4f}"
+            )
+            score = float(val_metrics[args.select_metric])
+            rank_tuple = (
+                score,
+                float(val_metrics.get("condition_top1_acc", 0.0)),
+                float(val_metrics.get("case_cellmean_top5_acc", 0.0)),
+                -float(val_metrics.get("loss", 0.0)),
+            )
+            if rank_tuple > best_rank_tuple:
+                best_score = score
+                best_rank_tuple = rank_tuple
+                best_epoch = epoch
+                best_state = {key: value.detach().cpu() for key, value in unwrapped_model.state_dict().items()}
+                best_val_metrics = dict(val_metrics)
+                no_improve = 0
+            else:
+                no_improve += 1
+        if is_distributed:
+            stop_flag = torch.tensor([0], device=device)
+            if is_main_process and args.early_stop_patience > 0 and no_improve >= args.early_stop_patience:
+                print_sys(f"Early stopping at epoch {epoch} (no improve for {no_improve} epochs)")
+                stop_flag[0] = 1
+            dist.broadcast(stop_flag, src=0)
+            if stop_flag.item() == 1:
+                break
+        else:
+            if args.early_stop_patience > 0 and no_improve >= args.early_stop_patience:
+                print_sys(f"Early stopping at epoch {epoch} (no improve for {no_improve} epochs)")
+                break
+        if is_distributed:
+            dist.barrier()
+
+    checkpoint_path = output_dir / "best_finetune_gene_inverse.pth"
+    checkpoint_saved = False
+    checkpoint_payload = None
+    if is_main_process:
+        if best_state is None:
+            raise RuntimeError("Training did not produce a valid checkpoint")
+        history_df = pd.DataFrame(history)
+        history_df.to_csv(output_dir / "train_history.csv", index=False)
+        checkpoint_payload = {
+            "model_state_dict": best_state,
+            "epoch": int(best_epoch),
+            "val_metrics": best_val_metrics,
+            "class_table": class_df.to_dict(orient="records"),
+            "pool_method": args.pool_method,
+            "head_hidden_dim": int(args.head_hidden_dim),
+            "head_dropout": float(args.head_dropout),
+            "head_residual": bool(args.head_residual),
+            "temperature": float(args.temperature),
+            "finetune_mode": args.finetune_mode,
+            "run_config": vars(args),
+        }
+        if args.save_best_checkpoint or is_distributed:
+            torch.save(checkpoint_payload, checkpoint_path)
+            checkpoint_saved = True
+            if args.save_best_checkpoint:
+                print_sys(f"Saved best checkpoint to {checkpoint_path}")
+            else:
+                print_sys(
+                    f"Saved temporary checkpoint to {checkpoint_path} "
+                    "(required for distributed test evaluation; file will be removed after evaluation)."
+                )
+
+    if is_distributed:
+        dist.barrier()
+        saved_checkpoint = torch.load(checkpoint_path, map_location="cpu")
+        unwrapped_model.load_state_dict(saved_checkpoint["model_state_dict"], strict=True)
+        dist.barrier()
+        if is_main_process and not args.save_best_checkpoint and checkpoint_path.exists():
+            checkpoint_path.unlink()
+            print_sys(f"Removed temporary checkpoint {checkpoint_path}")
+    else:
+        unwrapped_model.load_state_dict(best_state, strict=True)
+    test_metrics, test_logits, test_targets, test_row_ids = evaluate_loader(
+        model=unwrapped_model,
+        dataloader=test_loader,
+        device=device,
+        scfm_index_positions=scfm_index_positions,
+        loss_fn=eval_loss_fn,
+        return_outputs=True,
+        is_distributed=is_distributed,
+        world_size=world_size,
+        rank=rank,
+    )
+
+    if is_main_process:
+        test_condition_names = [test_row_to_condition[int(row)] for row in test_row_ids.tolist()]
+        test_condition_metrics, aggregated_logits, aggregated_targets, ordered_conditions = compute_condition_hit_rates(
+            logits=test_logits,
+            targets=test_targets,
+            condition_names=test_condition_names,
+        )
+        test_metrics.update(test_condition_metrics)
+        test_relevant_metrics = compute_condition_relevant_hit_rates(
+            aggregated_logits=aggregated_logits,
+            aggregated_targets=aggregated_targets,
+            id_to_condition=id_to_condition,
+        )
+        test_metrics.update(test_relevant_metrics)
+        test_case_cellmean_metrics = compute_case_cellmean_hit_rates(
+            logits=test_logits,
+            targets=test_targets,
+            condition_names=test_condition_names,
+            id_to_condition=id_to_condition,
+        )
+        test_metrics.update(test_case_cellmean_metrics)
+
+        if args.export_test_logits:
+            logits_df = pd.DataFrame(test_logits.numpy(), columns=id_to_condition)
+            logits_df.insert(0, "cell_name", adata.obs_names[test_row_ids].astype(str).tolist())
+            logits_df.insert(1, "row_index", test_row_ids.astype(int))
+            logits_path = output_dir / "test_logits.csv"
+            logits_df.to_csv(logits_path, index=False)
+            print_sys(f"Saved test logits to {logits_path}")
+
+        pred_ids = test_logits.argmax(dim=1).numpy()
+        true_ids = test_targets.numpy()
+        pred_df = pd.DataFrame(
+            {
+                "row_index": test_row_ids.astype(int),
+                "cell_name": adata.obs_names[test_row_ids].astype(str).tolist(),
+                "true_class_id": true_ids.astype(int),
+                "pred_class_id": pred_ids.astype(int),
+                "true_condition": [id_to_condition[idx] for idx in true_ids],
+                "pred_condition": [id_to_condition[idx] for idx in pred_ids],
+                "true_gene": [id_to_gene[idx] for idx in true_ids],
+                "pred_gene": [id_to_gene[idx] for idx in pred_ids],
+                "correct": (true_ids == pred_ids).astype(int),
+            }
+        )
+        pred_df.to_csv(output_dir / "test_predictions.csv", index=False)
+
+        condition_ranking_records = build_condition_ranking_records(
+            aggregated_logits=aggregated_logits,
+            aggregated_targets=aggregated_targets,
+            ordered_conditions=ordered_conditions,
+            id_to_condition=id_to_condition,
+            top_k=10,
+        )
+        with (output_dir / "test_condition_rankings.jsonl").open("w", encoding="utf-8") as handle:
+            for record in condition_ranking_records:
+                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+        summary = {
+            "algorithm_name": args.algorithm_name,
+            "protocol": args.protocol or None,
+            "best_epoch": int(best_epoch),
+            "selection_metric": args.select_metric,
+            "best_val_score": float(best_score),
+            "selection_tiebreaker": [
+                "condition_top1_acc",
+                "case_cellmean_top5_acc",
+                "lower_val_loss",
+            ],
+            "pool_method": args.pool_method,
+            "finetune_mode": args.finetune_mode,
+            "head_residual": bool(args.head_residual),
+            "checkpoint_saved": bool(args.save_best_checkpoint and checkpoint_saved),
+            "temperature": float(args.temperature),
+            "class_balance_power": float(args.class_balance_power),
+            "label_smoothing": float(args.label_smoothing),
+            "n_classes": int(len(class_df)),
+            "n_train_conditions": int(len(train_conditions)),
+            "n_val_conditions": int(len(val_conditions)),
+            "n_test_conditions": int(len(test_conditions)),
+            "n_train_cells": int(train_rows.size),
+            "n_val_cells": int(val_rows.size),
+            "n_test_cells": int(test_rows.size),
+            **{f"best_val_{k}": float(v) for k, v in best_val_metrics.items()},
+            **{f"test_{k}": float(v) for k, v in test_metrics.items()},
+            "alignment_notes": [
+                "Candidate universe includes selected single and double perturbations.",
+                "Top-k metrics are reported for k=1..10.",
+                "Primary comparison metric should use condition_topk_acc on perturbation-level aggregated test conditions.",
+                "scGPT-style relevant-hit metrics are exported as condition_relevant_topk_acc.",
+                "Paper-style non-integer hit rates can be approximated via case_cellmean_topk_acc metrics.",
+                "The exact CellFM 20-gene list is not published in code here; benchmark gene pool must be checked before claiming exact reproduction.",
+            ],
+        }
+        best_val_table = build_topk_table_rows(
+            metrics=best_val_metrics,
+            split_name="best_val",
+            algorithm_name=args.algorithm_name,
+        )
+        test_table = build_topk_table_rows(
+            metrics=test_metrics,
+            split_name="test",
+            algorithm_name=args.algorithm_name,
+        )
+        metrics_table_df = pd.concat([best_val_table, test_table], axis=0, ignore_index=True)
+        metrics_table_df.to_csv(output_dir / "metrics_table.csv", index=False)
+        print_topk_table(metrics_table_df, title="Top-k Metrics Table")
+
+        with (output_dir / "summary_metrics.json").open("w", encoding="utf-8") as handle:
+            json.dump(summary, handle, indent=2, ensure_ascii=False)
+        test_only_summary = {
+            "algorithm_name": args.algorithm_name,
+            "protocol": args.protocol or None,
+            "selection_metric": args.select_metric,
+            "best_epoch": int(best_epoch),
+            "pool_method": args.pool_method,
+            "finetune_mode": args.finetune_mode,
+            "head_residual": bool(args.head_residual),
+            "temperature": float(args.temperature),
+            "class_balance_power": float(args.class_balance_power),
+            "label_smoothing": float(args.label_smoothing),
+            "n_test_conditions": int(len(test_conditions)),
+            "n_test_cells": int(test_rows.size),
+            **{k: float(v) for k, v in test_metrics.items()},
+        }
+        with (output_dir / "test_summary_metrics.json").open("w", encoding="utf-8") as handle:
+            json.dump(test_only_summary, handle, indent=2, ensure_ascii=False)
+        with (output_dir / "run_config.json").open("w", encoding="utf-8") as handle:
+            json.dump(vars(args), handle, indent=2, ensure_ascii=False)
+        print_sys("Test Summary")
+        print_sys(json.dumps(test_only_summary, indent=2, ensure_ascii=False))
+
+    if is_distributed:
+        dist.barrier()
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

BioMod/gene_inverse_norman/protocols/cellfm/genes.txt

@@ -0,0 +1,20 @@
+CBFA2T3
+CEBPA
+CEBPE
+CNN1
+ETS2
+FEV
+FOXA1
+FOXA3
+FOXF1
+FOSB
+HOXB9
+IGDCC3
+IKZF3
+MAP2K3
+MAP2K6
+MAPK1
+PTPN12
+SNAI1
+TGFBR2
+UBASH3A

BioMod/gene_inverse_norman/protocols/cellfm/split.json

@@ -0,0 +1,59 @@
+{
+  "name": "cellfm",
+  "description": "CellFM-style reverse perturbation split on Norman (20 genes; 33 train = 16 doubles + 17 singles, 3 val doubles, 11 test doubles = 47 observed conditions).",
+  "include_singles": true,
+  "include_doubles": true,
+  "train": [
+    "CEBPA+CEBPE",
+    "CEBPE+ETS2",
+    "CEBPE+FOSB",
+    "FOXA3+FOXF1",
+    "FOXA1+HOXB9",
+    "FOXA3+HOXB9",
+    "ETS2+IKZF3",
+    "IKZF3+MAP2K3",
+    "IKZF3+MAP2K6",
+    "CNN1+MAPK1",
+    "IKZF3+MAPK1",
+    "CEBPE+PTPN12",
+    "PTPN12+SNAI1",
+    "ETS2+TGFBR2",
+    "IGDCC3+TGFBR2",
+    "MAPK1+TGFBR2",
+    "CBFA2T3",
+    "CEBPA",
+    "CEBPE",
+    "CNN1",
+    "ETS2",
+    "FEV",
+    "FOXA1",
+    "FOXF1",
+    "HOXB9",
+    "IGDCC3",
+    "IKZF3",
+    "MAP2K6",
+    "MAPK1",
+    "PTPN12",
+    "SNAI1",
+    "TGFBR2",
+    "UBASH3A"
+  ],
+  "val": [
+    "FOXA1+FOXF1",
+    "ETS2+MAPK1",
+    "FOSB+PTPN12"
+  ],
+  "test": [
+    "CEBPE+CNN1",
+    "CNN1+ETS2",
+    "CBFA2T3+FEV",
+    "FOXA1+FOXA3",
+    "FOXF1+HOXB9",
+    "ETS2+IGDCC3",
+    "FOSB+IKZF3",
+    "MAP2K3+MAP2K6",
+    "IGDCC3+MAPK1",
+    "CNN1+UBASH3A",
+    "PTPN12+UBASH3A"
+  ]
+}

BioMod/gene_inverse_norman/protocols/scgpt/genes.txt

@@ -0,0 +1,20 @@
+CEBPA
+CEBPE
+CNN1
+FEV
+FOSB
+FOXA1
+FOXF1
+HOXB9
+IGDCC3
+IKZF3
+KLF1
+MAP2K6
+MAPK1
+OSR2
+PRTG
+PTPN9
+SAMD1
+TGFBR2
+UBASH3B
+ZBTB25

BioMod/gene_inverse_norman/protocols/scgpt/split.json

@@ -0,0 +1,61 @@
+{
+  "name": "scgpt",
+  "description": "scGPT-style reverse perturbation split on Norman (20 genes; 39 train = 20 doubles + 19 singles, 3 val doubles, 7 test doubles = 49 observed conditions).",
+  "include_singles": true,
+  "include_doubles": true,
+  "train": [
+    "UBASH3B+ZBTB25",
+    "SAMD1+UBASH3B",
+    "SAMD1+TGFBR2",
+    "PRTG+TGFBR2",
+    "OSR2+UBASH3B",
+    "MAPK1+TGFBR2",
+    "KLF1+TGFBR2",
+    "KLF1+MAP2K6",
+    "IKZF3+MAPK1",
+    "IKZF3+MAP2K6",
+    "IGDCC3+TGFBR2",
+    "FOXF1+HOXB9",
+    "FOXA1+KLF1",
+    "FOXA1+FOXF1",
+    "FOSB+OSR2",
+    "CEBPE+KLF1",
+    "CEBPE+FOSB",
+    "CEBPE+CNN1",
+    "CEBPA+KLF1",
+    "CEBPA+CEBPE",
+    "ZBTB25",
+    "UBASH3B",
+    "TGFBR2",
+    "SAMD1",
+    "PTPN9",
+    "PRTG",
+    "OSR2",
+    "MAPK1",
+    "MAP2K6",
+    "KLF1",
+    "IKZF3",
+    "IGDCC3",
+    "HOXB9",
+    "FOXF1",
+    "FOXA1",
+    "FEV",
+    "CNN1",
+    "CEBPE",
+    "CEBPA"
+  ],
+  "val": [
+    "IGDCC3+ZBTB25",
+    "IGDCC3+PRTG",
+    "CNN1+UBASH3B"
+  ],
+  "test": [
+    "PTPN9+UBASH3B",
+    "MAPK1+PRTG",
+    "IGDCC3+MAPK1",
+    "FOXA1+HOXB9",
+    "FOSB+UBASH3B",
+    "FOSB+IKZF3",
+    "CNN1+MAPK1"
+  ]
+}

BioMod/gene_inverse_norman/utils.py

@@ -0,0 +1,166 @@
+import pickle
+import sys
+from copy import deepcopy
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Set, Union
+
+import pandas as pd
+import torch
+
+
+def print_sys(message: str) -> None:
+    print(message, flush=True, file=sys.stderr)
+
+
+def strip_gene_version(gene_id: str) -> str:
+    return gene_id.split(".", 1)[0]
+
+
+def canonicalize_target_gene_id(
+    raw_gene: object,
+    symbol_to_ensembl: Dict[str, str],
+    target_gene_set: Set[str],
+) -> Optional[str]:
+    if pd.isna(raw_gene):
+        return None
+
+    raw_gene_str = str(raw_gene).strip()
+    if not raw_gene_str:
+        return None
+
+    direct_candidates = [raw_gene_str, strip_gene_version(raw_gene_str)]
+    for candidate in direct_candidates:
+        if candidate in target_gene_set:
+            return candidate
+
+    mapped_gene = symbol_to_ensembl.get(raw_gene_str)
+    if mapped_gene is None:
+        mapped_gene = symbol_to_ensembl.get(strip_gene_version(raw_gene_str))
+
+    if mapped_gene is None or pd.isna(mapped_gene):
+        return None
+
+    mapped_gene_str = strip_gene_version(str(mapped_gene).strip())
+    return mapped_gene_str or None
+
+
+def resolve_input_gene_ids(
+    adata,
+    gene_col: str,
+    symbol_to_ensembl: Dict[str, str],
+    target_gene_set: Set[str],
+) -> List[str]:
+    if gene_col not in adata.var.columns:
+        raise KeyError(f"gene_col '{gene_col}' not found in adata.var.columns")
+
+    raw_genes = [str(gene) if not pd.isna(gene) else "" for gene in adata.var[gene_col].tolist()]
+    resolved_gene_ids: List[str] = []
+
+    direct_hits = 0
+    mapped_hits = 0
+    unresolved = 0
+
+    for index, raw_gene in enumerate(raw_genes):
+        resolved_gene = canonicalize_target_gene_id(raw_gene, symbol_to_ensembl, target_gene_set)
+        if resolved_gene is None:
+            resolved_gene = f"unknown{index}"
+            unresolved += 1
+        elif strip_gene_version(raw_gene) == resolved_gene:
+            direct_hits += 1
+        else:
+            mapped_hits += 1
+        resolved_gene_ids.append(resolved_gene)
+
+    print_sys(
+        "Gene ID resolution summary: "
+        f"direct_hits={direct_hits}, mapped_hits={mapped_hits}, unresolved={unresolved}"
+    )
+    return resolved_gene_ids
+
+
+def reindex_tensor_universal(
+    tensor: torch.Tensor,
+    index_positions: Sequence[int],
+    dim: int,
+    filler: str = "0",
+    device: Union[str, torch.device] = "cpu",
+) -> torch.Tensor:
+    index_tensor = torch.tensor(index_positions, device=device)
+    tensor_shape = list(tensor.shape)
+    tensor_shape[dim] = len(index_positions)
+    index_shape = [1] * len(tensor_shape)
+    index_shape[dim] = len(index_positions)
+    index_tensor = index_tensor.view(*index_shape).expand(*tensor_shape)
+
+    if filler == "0":
+        padder_shape = deepcopy(tensor_shape)
+        padder_shape[dim] = 1
+        padder = torch.zeros(padder_shape, dtype=tensor.dtype, device=device)
+    elif filler == "mean":
+        padder = tensor.mean(dim=dim, keepdim=True)
+    else:
+        raise ValueError("filler should be 0 or mean")
+
+    expanded_tensor = torch.cat((tensor, padder), dim=dim)
+    return torch.gather(expanded_tensor, dim, index_tensor)
+
+
+def parse_condition(condition: str) -> List[str]:
+    return [gene for gene in str(condition).split("+") if gene and gene != "ctrl"]
+
+
+def canonical_condition_name(condition: str) -> str:
+    genes = sorted(set(parse_condition(condition)))
+    if not genes:
+        return "ctrl"
+    return "+".join(genes)
+
+
+def project_root() -> Path:
+    return Path(__file__).resolve().parents[2]
+
+
+def default_repo_path(*parts: str) -> str:
+    return str(project_root().joinpath(*parts))
+
+
+def load_pretrained_backbone(
+    scfm_hyper_params_path: Union[str, Path],
+    scfm_ckpt_path: Union[str, Path],
+    gene2vec_path: Union[str, Path],
+    device: torch.device,
+):
+    repo_root = project_root()
+    if str(repo_root) not in sys.path:
+        sys.path.insert(0, str(repo_root))
+
+    from BioMod.models_performer_module import (
+        PerformerEncoderOnly_Module,
+        sanitize_performer_module_config,
+    )
+
+    with open(scfm_hyper_params_path, "rb") as handle:
+        scfm_hyper_params = pickle.load(handle)
+
+    scfm_hyper_params.pop("module_embed_chunk_padded_gene_budget", None)
+    scfm_hyper_params = sanitize_performer_module_config(
+        scfm_hyper_params,
+        source=str(scfm_hyper_params_path),
+        log_fn=print_sys,
+    )
+    scfm_hyper_params["gene2vec_file"] = str(gene2vec_path)
+
+    model = PerformerEncoderOnly_Module(**scfm_hyper_params).to(device)
+    ckpt = torch.load(scfm_ckpt_path, map_location=device)
+    model_state_dict = ckpt.get("model_state_dict", ckpt)
+    filtered_state_dict = {
+        key: value
+        for key, value in model_state_dict.items()
+        if "go_conv" not in key and "G_go" not in key.lower()
+    }
+    model.load_state_dict(filtered_state_dict, strict=True)
+    for parameter in model.parameters():
+        parameter.requires_grad = False
+    model.eval()
+
+    return model, scfm_hyper_params

BioMod/gene_perturbation/README.md

@@ -0,0 +1,265 @@
+# Adapting PerformerEncoderOnly_Module to Gene Perturbation
+
+This document explains how to adapt the PerformerEncoderOnly_Module model to the
+gene perturbation task (GEARS).
+
+## Overview
+
+This adaptation replaces the original GEARS gene perturbation model based on
+`DualEncoderSCFM` with `PerformerEncoderOnly_Module`. The main adaptation points
+are:
+
+1. **Model architecture adaptation**: replace the original dual-encoder model with a module-based encoder-only model.
+2. **Data format adaptation**: handle module-based input data.
+3. **Embedding extraction adaptation**: recover gene-level embeddings from module encodings.
+
+## Main Adaptations
+
+### 1. `SCFM_Module` (`model_module.py`)
+
+`SCFM_Module` wraps `PerformerEncoderOnly_Module` and provides these functions:
+
+- **Input handling**: reindex input gene expression values to the scfm gene order.
+- **Module encoding**: obtain module-level encodings from `PerformerEncoderOnly_Module`.
+- **Gene embedding recovery**: recover gene-level embeddings from module encodings.
+  - For each gene, find all modules that contain it.
+  - Average the encodings of those modules to obtain the gene embedding.
+  - If a gene is not included in any module, use the global average as a fallback.
+- **Output reindexing**: reindex gene embeddings back to the input gene order.
+
+**Key adaptation**:
+
+```python
+# Recover gene-level embeddings from module encodings.
+for gene_idx in range(len(self.scfm_index_positions) - 1):
+    if gene_idx in self.gene_to_modules:
+        module_indices = self.gene_to_modules[gene_idx]
+        module_embs = module_encodings[:, module_indices, :]  # B, num_modules_for_gene, D
+        gene_emb = module_embs.mean(dim=1)  # B, D
+    else:
+        gene_emb = module_encodings.mean(dim=1)  # B, D (global average)
+```
+
+### 2. `GEARS_Model_Module` (`model_module.py`)
+
+`GEARS_Model_Module` inherits the original GEARS architecture. Main changes:
+
+- **Use `SCFM_Module`** instead of the original `SCFM_different_genes`.
+- **Module data loading** supports loading module data from hyper-parameters or file paths.
+- **Other components remain unchanged**, including the GNN layers and decoder.
+
+### 3. Training Script (`train_module.py`)
+
+The training script is based on `GEARS/train_dual_ddp_reduce_dec.py`. Main changes:
+
+- **Model class**: use `GEARS_Model_Module` instead of `GEARS_Model_Acc`.
+- **New arguments**:
+  - `--module_data_path`: path to one module JSON file.
+  - `--module_data_dir`: directory containing module JSON files.
+  - `--gene_mapping_path`: CSV file mapping gene symbols to Ensembl IDs.
+
+## Usage
+
+### 1. Prepare Data
+
+Make sure the following files exist:
+
+- Perturbation data, such as Norman, Adamson, or Dixit.
+- scfm model checkpoint and hyper-parameter file.
+- Module data in JSON format.
+- Gene mapping file, such as `gene2ens.csv`.
+
+### 2. Run Training
+
+```bash
+# DDP training on multiple GPUs.
+python -m torch.distributed.launch --nproc_per_node=4 \
+    train_module.py \
+    --data_dir /path/to/data \
+    --data_name norman \
+    --scfm_hyper_params_path /path/to/hyper_params.pkl \
+    --scfm_ckpt_path /path/to/checkpoint.pth \
+    --module_data_dir /path/to/module_data \
+    --gene_mapping_path /path/to/gene2ens.csv \
+    --scfm_genes_list_path /path/to/selected_genes_27k.txt \
+    --scfm_gene2vec_file /path/to/selected_gene2vec_27k.npy \
+    --result_dir ./results_module \
+    --run_name norman_module_experiment \
+    --epochs 20 \
+    --batch_size 32 \
+    --lr 1e-3
+```
+
+### 3. Arguments
+
+**Required arguments**:
+
+- `--scfm_hyper_params_path`: path to the scfm model hyper-parameter file.
+- `--scfm_ckpt_path`: path to the scfm model checkpoint.
+- `--module_data_path` or `--module_data_dir`: module data path; provide one of them.
+
+**Recommended arguments**:
+
+- `--gene_mapping_path`: gene mapping file, especially when module data use gene symbols.
+- `--scfm_genes_list_path`: scfm gene list file.
+- `--scfm_gene2vec_file`: path to the gene2vec file.
+
+**Optional arguments**:
+
+- `--use_module_level`: use module-level features. This is experimental and is not fully supported at the moment.
+  - Default: False, meaning gene-level features are used and recommended.
+  - See `METHODS_COMPARISON.md` for details.
+
+## Comparison with Original sclong
+
+### Main Differences
+
+1. **Model architecture**:
+   - **Original sclong**: uses `DualEncoderSCFM` and directly processes gene sequences.
+   - **Module version**: uses `PerformerEncoderOnly_Module` and processes module-based data.
+
+2. **Data representation**:
+   - **Original sclong**: gene-level embeddings come directly from encoder outputs.
+   - **Module version**: gene-level embeddings are recovered from module encodings by mean pooling.
+
+3. **Input format**:
+   - **Original sclong**: directly inputs gene expression vectors.
+   - **Module version**: keeps the same external input format but internally processes data with module embeddings.
+
+### Performance Comparison
+
+To compare model performance:
+
+1. **Run original sclong**:
+
+   ```bash
+   cd GEARS
+   python train_dual_ddp_reduce_dec.py [arguments...]
+   ```
+
+2. **Run the module version**:
+
+   ```bash
+   cd performer_module/gene_perturbation
+   python train_module.py [arguments...]
+   ```
+
+3. **Compare metrics**:
+
+   - MSE (Mean Squared Error)
+   - Pearson correlation
+   - MSE and Pearson correlation for the top 20 DE genes
+   - Direction accuracy and related metrics
+
+Results are saved under each `result_dir`. Compare these metrics to evaluate the
+module version.
+
+## Technical Details
+
+### Module-to-Gene Mapping
+
+In `SCFM_Module`, `gene_to_modules` records which modules contain each gene:
+
+```python
+self.gene_to_modules = {}  # {gene_idx: [module_idx1, module_idx2, ...]}
+for module_idx, module_genes in enumerate(module_sets):
+    for gene_idx in module_genes:
+        if gene_idx < len(scfm_genes):
+            if gene_idx not in self.gene_to_modules:
+                self.gene_to_modules[gene_idx] = []
+            self.gene_to_modules[gene_idx].append(module_idx)
+```
+
+### Gene Embedding Recovery Strategy
+
+For each gene:
+
+1. If the gene appears in some modules, average the encodings of those modules.
+2. If the gene does not appear in any module, use the global average of all modules as a fallback.
+
+This strategy:
+
+- Makes full use of module information.
+- Provides a reasonable embedding for genes absent from all modules.
+- Keeps embedding dimensionality similar to the original model.
+
+## File Structure
+
+```text
+BioMod/gene_perturbation/                                # Perturbation response prediction: processing, training and inference
+├── data_utils_local.py
+├── gears_reduce_local.py
+├── inference_local.py
+├── model_module.py
+├── pertdata_local.py
+├── README.md
+├── scfm_utils_local.py
+├── train_module.py
+└── utils.py
+```
+
+## Notes
+
+1. **Module data format**: make sure the gene indices or names in the module data match the scfm gene list.
+2. **Gene mapping**: provide `gene_mapping_path` if the module data use gene symbols.
+3. **Memory usage**: the module version may require more memory because it maintains the module-to-gene mapping.
+4. **Training stability**: use the same learning rate and training strategy as original sclong when possible.
+
+## Troubleshooting
+
+### Issue 1: Module Data Not Found
+
+- Check whether `--module_data_path` or `--module_data_dir` is correct.
+- Make sure the module data file exists and has the expected format.
+
+### Issue 2: Gene Mapping Failed
+
+- Check whether `--gene_mapping_path` is correct.
+- Make sure the gene mapping file is a CSV containing `Approved symbol` and `Ensembl gene ID` columns.
+
+### Issue 3: Model Loading Failed
+
+- Check whether the checkpoint path is correct.
+- Make sure `hyper_params` contains the required module information.
+
+## Gene-Based vs Module-Based Methods
+
+This implementation supports two feature extraction modes controlled by
+`--use_module_level`.
+
+### Method 1: Gene-Based (Default, Recommended)
+
+```bash
+# Default mode; do not pass --use_module_level or set it to False.
+python train_module.py [other arguments...]
+```
+
+**Characteristics**:
+
+- Recovers gene-level embeddings from module encodings.
+- Fully compatible with the original GEARS architecture.
+- Downstream tasks operate at gene level, including the GNN and decoder.
+- Outputs gene-level predictions.
+
+### Method 2: Module-Based (Experimental, Not Fully Implemented)
+
+```bash
+# Enable module-level mode.
+python train_module.py --use_module_level [other arguments...]
+```
+
+**Characteristics**:
+
+- Directly uses module-level encodings.
+- Preserves the full module-level information.
+- **Note**: module-level downstream tasks are not fully implemented in the current version.
+- Requires redesigning the downstream architecture, including the GNN and decoder.
+
+See `METHODS_COMPARISON.md` for the detailed comparison and rationale.
+
+## References
+
+- Original GEARS code: `GEARS/gears_reduce/`
+- PerformerEncoderOnly_Module model: `performer_module/models_performer_module.py`
+- Module data processing: `performer_module/module_dataset.py`
+- Method comparison: `METHODS_COMPARISON.md`

BioMod/gene_perturbation/data_utils_local.py

@@ -0,0 +1,501 @@
+"""
+================================================================================
+Data processing utility functions (data_utils_local.py)
+================================================================================
+
+[Module function]
+Provides tool functions for genetic perturbation data preprocessing, including:
+1. Calculation of differentially expressed genes (DE genes)
+2. Data set partitioning strategy
+3. Dropout gene identification
+
+[Key concepts]
+- Differentially expressed genes: genes with the most significant expression changes after perturbation, used to evaluate model prediction quality
+- Dropout: Common technical noise in single-cell sequencing, some genes cannot be detected in some cells
+- Simulation split: simulate the data division of real scenarios and test the model’s generalization ability to unseen disturbances
+================================================================================
+"""
+
+import pandas as pd
+import numpy as np
+import scanpy as sc
+from random import shuffle
+sc.settings.verbosity = 0
+from tqdm import tqdm
+import requests
+import os, sys
+
+import warnings
+warnings.filterwarnings("ignore")
+
+from utils import parse_single_pert, parse_combo_pert, parse_any_pert, print_sys
+
+def rank_genes_groups_by_cov(
+    adata,
+    groupby,
+    control_group,
+    covariate,
+    pool_doses=False,
+    n_genes=50,
+    rankby_abs=True,
+    key_added='rank_genes_groups_cov',
+    return_dict=False,
+):
+    """
+    Calculate differentially expressed genes grouped by covariates
+
+    [Function]
+    For each cell type (covariate), the differentially expressed genes of each perturbation condition relative to the control group were calculated.
+
+    [Parameter description]
+    - groupby: grouping variable, usually disturbance condition
+    - control_group: control group name
+    - covariate: covariate, usually cell type
+    - rankby_abs: Sort by absolute value (both up-regulated and down-regulated genes are considered)
+
+    [Output]
+    The result is saved in adata.uns[key_added]
+    """
+    gene_dict = {}
+    cov_categories = adata.obs[covariate].unique()
+    for cov_cat in cov_categories:
+        #name of the control group in the groupby obs column
+        control_group_cov = '_'.join([cov_cat, control_group])
+
+        #subset adata to cells belonging to a covariate category
+        adata_cov = adata[adata.obs[covariate]==cov_cat]
+
+        #compute DEGs
+        sc.tl.rank_genes_groups(
+            adata_cov,
+            groupby=groupby,
+            reference=control_group_cov,
+            rankby_abs=rankby_abs,
+            n_genes=n_genes,
+            use_raw=False
+        )
+
+        #add entries to dictionary of gene sets
+        de_genes = pd.DataFrame(adata_cov.uns['rank_genes_groups']['names'])
+        for group in de_genes:
+            gene_dict[group] = de_genes[group].tolist()
+
+    adata.uns[key_added] = gene_dict
+
+    if return_dict:
+        return gene_dict
+
+
+def get_DE_genes(adata):
+    """
+    Calculate differentially expressed genes
+
+    [Function]
+    Differentially expressed genes relative to the control group were calculated for each perturbation condition.
+    This is one of the core indicators of GEARS evaluation: whether the DE genes predicted by the model are consistent with reality.
+
+    [Processing Process]
+    1. Construct the name of the disturbance condition (cell_type + condition + dose)
+    2. Calculate DE genes grouped by cell type
+    """
+    adata.obs.loc[:, 'dose_val'] = adata.obs.condition.apply(lambda x: '1+1' if len(x.split('+')) == 2 else '1')
+    adata.obs.loc[:, 'control'] = adata.obs.condition.apply(lambda x: 0 if len(x.split('+')) == 2 else 1)
+    adata.obs.loc[:, 'condition_name'] =  adata.obs.apply(lambda x: '_'.join([x.cell_type, x.condition, x.dose_val]), axis = 1)
+    rank_genes_groups_by_cov(adata,
+                         groupby='condition_name',
+                         covariate='cell_type',
+                         control_group='ctrl_1',
+                         n_genes=len(adata.var),
+                         key_added = 'rank_genes_groups_cov_all')
+    return adata
+
+def get_dropout_non_zero_genes(adata):
+    """
+    Identify dropout and non-zero expression genes
+
+    [Function]
+    There are a large number of dropouts (genes that should be expressed but are not detected) in single-cell sequencing.
+    This function recognizes:
+    1. Non-zero expressed genes: genes detected in perturbed cells
+    2. Non-dropout genes: genes that are 0 in both the perturbation and control groups (may be truly not expressed)
+
+    [Why is it important]
+    - Dropout genes should be excluded from the evaluation as they introduce unpredictable noise
+    - top_non_dropout_de_20: used to evaluate the prediction ability of the model on reliable genes
+    """
+    # calculate mean expression for each condition
+    unique_conditions = adata.obs.condition.unique()
+    conditions2index = {}
+    for i in unique_conditions:
+        conditions2index[i] = np.where(adata.obs.condition == i)[0]
+
+    condition2mean_expression = {}
+    for i, j in conditions2index.items():
+        condition2mean_expression[i] = np.mean(adata.X[j], axis = 0)
+    pert_list = np.array(list(condition2mean_expression.keys()))
+    mean_expression = np.array(list(condition2mean_expression.values())).reshape(len(adata.obs.condition.unique()), adata.X.toarray().shape[1])
+    ctrl = mean_expression[np.where(pert_list == 'ctrl')[0]]
+
+    ## in silico modeling and upperbounding
+    pert2pert_full_id = dict(adata.obs[['condition', 'condition_name']].values)
+    pert_full_id2pert = dict(adata.obs[['condition_name', 'condition']].values)
+
+    gene_id2idx = dict(zip(adata.var.index.values, range(len(adata.var))))
+    gene_idx2id = dict(zip(range(len(adata.var)), adata.var.index.values))
+
+    non_zeros_gene_idx = {}
+    top_non_dropout_de_20 = {}
+    top_non_zero_de_20 = {}
+    non_dropout_gene_idx = {}
+
+    for pert in adata.uns['rank_genes_groups_cov_all'].keys():
+        p = pert_full_id2pert[pert]
+        X = np.mean(adata[adata.obs.condition == p].X, axis = 0)
+
+        # Identify non-zero and zero expressed genes
+        non_zero = np.where(np.array(X)[0] != 0)[0]
+        zero = np.where(np.array(X)[0] == 0)[0]
+        # True dropout: both disturbance and control groups are 0
+        true_zeros = np.intersect1d(zero, np.where(np.array(ctrl)[0] == 0)[0])
+        non_dropouts = np.concatenate((non_zero, true_zeros))
+
+        # Get non-dropout genes among top DE genes
+        top = adata.uns['rank_genes_groups_cov_all'][pert]
+        gene_idx_top = [gene_id2idx[i] for i in top]
+
+        non_dropout_20 = [i for i in gene_idx_top if i in non_dropouts][:20]
+        non_dropout_20_gene_id = [gene_idx2id[i] for i in non_dropout_20]
+
+        non_zero_20 = [i for i in gene_idx_top if i in non_zero][:20]
+        non_zero_20_gene_id = [gene_idx2id[i] for i in non_zero_20]
+
+        non_zeros_gene_idx[pert] = np.sort(non_zero)
+        non_dropout_gene_idx[pert] = np.sort(non_dropouts)
+        top_non_dropout_de_20[pert] = np.array(non_dropout_20_gene_id)
+        top_non_zero_de_20[pert] = np.array(non_zero_20_gene_id)
+
+    non_zero = np.where(np.array(X)[0] != 0)[0]
+    zero = np.where(np.array(X)[0] == 0)[0]
+    true_zeros = np.intersect1d(zero, np.where(np.array(ctrl)[0] == 0)[0])
+    non_dropouts = np.concatenate((non_zero, true_zeros))
+
+    adata.uns['top_non_dropout_de_20'] = top_non_dropout_de_20
+    adata.uns['non_dropout_gene_idx'] = non_dropout_gene_idx
+    adata.uns['non_zeros_gene_idx'] = non_zeros_gene_idx
+    adata.uns['top_non_zero_de_20'] = top_non_zero_de_20
+
+    return adata
+
+
+class DataSplitter():
+    """
+    Data classification class
+
+    [Function]
+    Handle training/validation/test partitioning of genetically perturbed datasets.
+    Different partitioning strategies evaluate the capabilities of different aspects of the model.
+
+    [Detailed explanation of partition strategy]
+    - simulation: the most comprehensive division, the test set includes:
+        * unseen_single: Unseen single gene perturbation
+        * combo_seen0: a combination that neither gene has seen
+        * combo_seen1: the combination seen by one of the genes
+        * combo_seen2: Both genes have been seen but the combination has not been seen
+    """
+    def __init__(self, adata, split_type='single', seen=0):
+        self.adata = adata
+        self.split_type = split_type
+        self.seen = seen
+
+    def split_data(self, test_size=0.1, test_pert_genes=None,
+                   test_perts=None, split_name='split', seed=None, val_size = 0.1,
+                   train_gene_set_size = 0.75, combo_seen2_train_frac = 0.75, only_test_set_perts = False):
+        """
+        Perform data partitioning
+
+        [Key parameters]
+        - train_gene_set_size: training set gene proportion
+        - combo_seen2_train_frac: The proportion of combo_seen2 used for training
+        """
+        np.random.seed(seed=seed)
+        unique_perts = [p for p in self.adata.obs['condition'].unique() if
+                        p != 'ctrl']
+
+        if self.split_type == 'simulation':
+            train, test, test_subgroup = self.get_simulation_split(unique_perts,
+                                                                  train_gene_set_size,
+                                                                  combo_seen2_train_frac,
+                                                                  seed, test_perts, only_test_set_perts)
+            train, val, val_subgroup = self.get_simulation_split(train,
+                                                                  0.9,
+                                                                  0.9,
+                                                                  seed)
+            ## adding back ctrl to train...
+            train.append('ctrl')
+        elif self.split_type == 'simulation_single':
+            train, test, test_subgroup = self.get_simulation_split_single(unique_perts,
+                                                                  train_gene_set_size,
+                                                                  seed, test_perts, only_test_set_perts)
+            train, val, val_subgroup = self.get_simulation_split_single(train,
+                                                                  0.9,
+                                                                  seed)
+        elif self.split_type == 'no_test':
+            print('test_pert_genes',str(test_pert_genes))
+            print('test_perts',str(test_perts))
+
+            train, val = self.get_split_list(unique_perts,
+                                          test_pert_genes=test_pert_genes,
+                                          test_perts=test_perts,
+                                          test_size=test_size)
+        else:
+            train, test = self.get_split_list(unique_perts,
+                                          test_pert_genes=test_pert_genes,
+                                          test_perts=test_perts,
+                                          test_size=test_size)
+
+            train, val = self.get_split_list(train, test_size=val_size)
+
+        map_dict = {x: 'train' for x in train}
+        map_dict.update({x: 'val' for x in val})
+        if self.split_type != 'no_test':
+            map_dict.update({x: 'test' for x in test})
+        map_dict.update({'ctrl': 'train'})
+
+        self.adata.obs[split_name] = self.adata.obs['condition'].map(map_dict)
+
+        if self.split_type == 'simulation':
+            return self.adata, {'test_subgroup': test_subgroup,
+                                'val_subgroup': val_subgroup
+                               }
+        else:
+            return self.adata
+    
+    def get_simulation_split_single(self, pert_list, train_gene_set_size = 0.85, seed = 1, test_set_perts = None, only_test_set_perts = False):
+        unique_pert_genes = self.get_genes_from_perts(pert_list)
+        
+        pert_train = []
+        pert_test = []
+        np.random.seed(seed=seed)
+        
+        if only_test_set_perts and (test_set_perts is not None):
+            ood_genes = np.array(test_set_perts)
+            train_gene_candidates = np.setdiff1d(unique_pert_genes, ood_genes)
+        else:
+            ## a pre-specified list of genes
+            train_gene_candidates = np.random.choice(unique_pert_genes,
+                                                    int(len(unique_pert_genes) * train_gene_set_size), replace = False)
+
+            if test_set_perts is not None:
+                num_overlap = len(np.intersect1d(train_gene_candidates, test_set_perts))
+                train_gene_candidates = train_gene_candidates[~np.isin(train_gene_candidates, test_set_perts)]
+                ood_genes_exclude_test_set = np.setdiff1d(unique_pert_genes, np.union1d(train_gene_candidates, test_set_perts))
+                train_set_addition = np.random.choice(ood_genes_exclude_test_set, num_overlap, replace = False)
+                train_gene_candidates = np.concatenate((train_gene_candidates, train_set_addition))
+                
+            ## ood genes
+            ood_genes = np.setdiff1d(unique_pert_genes, train_gene_candidates)  
+        
+        pert_single_train = self.get_perts_from_genes(train_gene_candidates, pert_list,'single')
+        unseen_single = self.get_perts_from_genes(ood_genes, pert_list, 'single')
+        assert len(unseen_single) + len(pert_single_train) == len(pert_list)
+        
+        return pert_single_train, unseen_single, {'unseen_single': unseen_single}
+    
+    def get_simulation_split(self, pert_list, train_gene_set_size = 0.85, combo_seen2_train_frac = 0.85, seed = 1, test_set_perts = None, only_test_set_perts = False):
+        """
+        Create data partitions that simulate real-world scenarios
+
+        [Division strategy]
+        This is the most important evaluation method in the GEARS paper, testing the model's generalization ability to unseen perturbations.
+
+        The test set contains the following subtypes:
+        1. unseen_single: Single gene perturbation that has never been seen during training
+        2. combo_seen0: Combination perturbation that neither gene has seen
+        3. combo_seen1: Combination perturbation seen by one of the genes
+        4. combo_seen2: Both genes have been seen, but this specific combination has not been seen
+
+        [Assessment significance]
+        - unseen_single: Evaluate the model's predictive ability for novel genetic perturbations
+        - combo_seen0: Evaluate model predictions for completely unknown combinations (most difficult)
+        - combo_seen1: Evaluate whether the model can infer combination effects from single gene knowledge
+        - combo_seen2: Evaluate whether the model can infer new combinations from known combinations
+        """
+        unique_pert_genes = self.get_genes_from_perts(pert_list)
+
+        pert_train = []
+        pert_test = []
+        np.random.seed(seed=seed)
+
+        if only_test_set_perts and (test_set_perts is not None):
+            ood_genes = np.array(test_set_perts)
+            train_gene_candidates = np.setdiff1d(unique_pert_genes, ood_genes)
+        else:
+            # Randomly select genes from the training set
+            train_gene_candidates = np.random.choice(unique_pert_genes,
+                                                    int(len(unique_pert_genes) * train_gene_set_size), replace = False)
+
+            if test_set_perts is not None:
+                num_overlap = len(np.intersect1d(train_gene_candidates, test_set_perts))
+                train_gene_candidates = train_gene_candidates[~np.isin(train_gene_candidates, test_set_perts)]
+                ood_genes_exclude_test_set = np.setdiff1d(unique_pert_genes, np.union1d(train_gene_candidates, test_set_perts))
+                train_set_addition = np.random.choice(ood_genes_exclude_test_set, num_overlap, replace = False)
+                train_gene_candidates = np.concatenate((train_gene_candidates, train_set_addition))
+
+            # OOD genes: genes not in the training set
+            ood_genes = np.setdiff1d(unique_pert_genes, train_gene_candidates)
+
+        # Obtain single gene perturbations of the training set
+        pert_single_train = self.get_perts_from_genes(train_gene_candidates, pert_list,'single')
+        pert_combo = self.get_perts_from_genes(train_gene_candidates, pert_list,'combo')
+        pert_train.extend(pert_single_train)
+
+        # ============================================================
+        # [Key] Divide test subtypes of combined disturbances
+        # ============================================================
+        # combo_seen1: One of the genes is in the training set
+        combo_seen1 = [x for x in pert_combo if len([t for t in x.split('+') if
+                                     t in train_gene_candidates]) == 1]
+        pert_test.extend(combo_seen1)
+
+        pert_combo = np.setdiff1d(pert_combo, combo_seen1)
+        # combo_seen2: Both genes are in the training set, but this combination is not in the training set
+        np.random.seed(seed=seed)
+        pert_combo_train = np.random.choice(pert_combo, int(len(pert_combo) * combo_seen2_train_frac), replace = False)
+
+        combo_seen2 = np.setdiff1d(pert_combo, pert_combo_train).tolist()
+        pert_test.extend(combo_seen2)
+        pert_train.extend(pert_combo_train)
+
+        # unseen_single: Unseen single gene perturbation
+        unseen_single = self.get_perts_from_genes(ood_genes, pert_list, 'single')
+        combo_ood = self.get_perts_from_genes(ood_genes, pert_list, 'combo')
+        pert_test.extend(unseen_single)
+
+        # combo_seen0: Neither gene has been seen
+        combo_seen0 = [x for x in combo_ood if len([t for t in x.split('+') if
+                                     t in train_gene_candidates]) == 0]
+        pert_test.extend(combo_seen0)
+
+        # Verify partition integrity
+        assert len(combo_seen1) + len(combo_seen0) + len(unseen_single) + len(pert_train) + len(combo_seen2) == len(pert_list)
+
+        return pert_train, pert_test, {'combo_seen0': combo_seen0,
+                                       'combo_seen1': combo_seen1,
+                                       'combo_seen2': combo_seen2,
+                                       'unseen_single': unseen_single}
+        
+    def get_split_list(self, pert_list, test_size=0.1,
+                       test_pert_genes=None, test_perts=None,
+                       hold_outs=True):
+        """
+        Splits a given perturbation list into train and test with no shared
+        perturbations
+        """
+
+        single_perts = [p for p in pert_list if 'ctrl' in p and p != 'ctrl']
+        combo_perts = [p for p in pert_list if 'ctrl' not in p]
+        unique_pert_genes = self.get_genes_from_perts(pert_list)
+        hold_out = []
+
+        if test_pert_genes is None:
+            test_pert_genes = np.random.choice(unique_pert_genes,
+                                        int(len(single_perts) * test_size))
+
+        # Only single unseen genes (in test set)
+        # Train contains both single and combos
+        if self.split_type == 'single' or self.split_type == 'single_only':
+            test_perts = self.get_perts_from_genes(test_pert_genes, pert_list,
+                                                   'single')
+            if self.split_type == 'single_only':
+                # Discard all combos
+                hold_out = combo_perts
+            else:
+                # Discard only those combos which contain test genes
+                hold_out = self.get_perts_from_genes(test_pert_genes, pert_list,
+                                                     'combo')
+
+        elif self.split_type == 'combo':
+            if self.seen == 0:
+                # NOTE: This can reduce the dataset size!
+                # To prevent this set 'holdouts' to False, this will cause
+                # the test set to have some perturbations with 1 gene seen
+                single_perts = self.get_perts_from_genes(test_pert_genes,
+                                                         pert_list, 'single')
+                combo_perts = self.get_perts_from_genes(test_pert_genes,
+                                                        pert_list, 'combo')
+
+                if hold_outs:
+                    # This just checks that none of the combos have 2 seen genes
+                    hold_out = [t for t in combo_perts if
+                                len([t for t in t.split('+') if
+                                     t not in test_pert_genes]) > 0]
+                combo_perts = [c for c in combo_perts if c not in hold_out]
+                test_perts = single_perts + combo_perts
+
+            elif self.seen == 1:
+                # NOTE: This can reduce the dataset size!
+                # To prevent this set 'holdouts' to False, this will cause
+                # the test set to have some perturbations with 2 genes seen
+                single_perts = self.get_perts_from_genes(test_pert_genes,
+                                                         pert_list, 'single')
+                combo_perts = self.get_perts_from_genes(test_pert_genes,
+                                                        pert_list, 'combo')
+
+                if hold_outs:
+                    # This just checks that none of the combos have 2 seen genes
+                    hold_out = [t for t in combo_perts if
+                                len([t for t in t.split('+') if
+                                     t not in test_pert_genes]) > 1]
+                combo_perts = [c for c in combo_perts if c not in hold_out]
+                test_perts = single_perts + combo_perts
+
+            elif self.seen == 2:
+                if test_perts is None:
+                    test_perts = np.random.choice(combo_perts,
+                                     int(len(combo_perts) * test_size))       
+                else:
+                    test_perts = np.array(test_perts)
+        else:
+            if test_perts is None:
+                test_perts = np.random.choice(combo_perts,
+                                    int(len(combo_perts) * test_size))
+        
+        train_perts = [p for p in pert_list if (p not in test_perts)
+                                        and (p not in hold_out)]
+        return train_perts, test_perts
+
+    def get_perts_from_genes(self, genes, pert_list, type_='both'):
+        """
+        Returns all single/combo/both perturbations that include a gene
+        """
+
+        single_perts = [p for p in pert_list if ('ctrl' in p) and (p != 'ctrl')]
+        combo_perts = [p for p in pert_list if 'ctrl' not in p]
+        
+        perts = []
+        
+        if type_ == 'single':
+            pert_candidate_list = single_perts
+        elif type_ == 'combo':
+            pert_candidate_list = combo_perts
+        elif type_ == 'both':
+            pert_candidate_list = pert_list
+            
+        for p in pert_candidate_list:
+            for g in genes:
+                if g in parse_any_pert(p):
+                    perts.append(p)
+                    break
+        return perts
+
+    def get_genes_from_perts(self, perts):
+        """
+        Returns list of genes involved in a given perturbation list
+        """
+
+        if type(perts) is str:
+            perts = [perts]
+        gene_list = [p.split('+') for p in np.unique(perts)]
+        gene_list = [item for sublist in gene_list for item in sublist]
+        gene_list = [g for g in gene_list if g != 'ctrl']
+        return np.unique(gene_list)

BioMod/gene_perturbation/gears_reduce_local.py

@@ -0,0 +1,761 @@
+"""
+================================================================================
+GEARS training and evaluation controller (gears_reduce_local.py)
+================================================================================
+
+[Module function]
+This is the core controller of the genetic perturbation mission, responsible for:
+1. Data loader configuration
+2. Model initialization and hyperparameter management
+3. Training loop and validation
+4. Test evaluation and result analysis
+
+[Reading suggestions]
+1. First look at `__init__`: understand the data loader configuration and key variable initialization
+2. Look again at `model_initialize`: decision graph construction, model classes, hyperparameters and freezing strategies
+3. Then look at `train`: determining optimization goals, gradient accumulation, verification indicators and optimal model selection criteria
+4. Finally, look at the test section: determine the final report indicators and simulation subset analysis caliber
+
+[Key parts that affect the results]
+- dict_filter: training loss is only calculated on non-zero genes
+- direction_lambda: direction consistency loss weight
+- finetune_method: freeze vs full parameter fine-tuning
+- G_coexpress/G_go: Genetic relationship graph construction parameters
+- accumulation_steps: equivalent batch size of gradient accumulation influence
+================================================================================
+"""
+
+from copy import deepcopy
+import argparse
+from time import time
+import sys, os
+import pickle
+
+import scanpy as sc
+import numpy as np
+
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.optim.lr_scheduler import StepLR
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+from torch_geometric.data import DataLoader
+
+# ============================================================
+# Use native modules (optimized versions)
+# ============================================================
+from model_module import GEARS_Model_Module
+from inference_local import evaluate, compute_metrics, deeper_analysis, \
+                  non_dropout_analysis, compute_synergy_loss
+from utils import loss_fct, uncertainty_loss_fct, parse_any_pert, \
+                  get_similarity_network, print_sys, GeneSimNetwork, \
+                  create_cell_graph_dataset_for_prediction, get_mean_control, \
+                  get_GI_genes_idx, get_GI_params
+from scfm_utils_local import *
+
+torch.manual_seed(0)
+
+import warnings
+warnings.filterwarnings("ignore")
+
+
+class GEARS:
+    """
+    GEARS training and evaluation controller
+
+    [Core functions]
+    1. Manage the entire training process (data, model, optimization, evaluation)
+    2. Support distributed training (DDP)
+    3. Support Weights & Biases experiment tracking
+    4. Support model saving/loading
+
+    [Usage process]
+    ```python
+    gears = GEARS(pert_data, local_rank, is_master, world_size, ...)
+    gears.model_initialize(hidden_size=64, model_class=GEARS_Model_Module, ...)
+    gears.train(epochs=20, result_dir='./results', lr=1e-3)
+    ```
+    """
+
+    def __init__(self, pert_data,
+                 local_rank,
+                 is_master: bool,
+                 world_size,
+                 train_bs = 48,
+                 test_bs = 48,
+                 device = 'cuda',
+                 weight_bias_track = True,
+                 proj_name = 'GEARS',
+                 exp_name = 'GEARS',):
+        """
+        Initialize the GEARS controller
+
+        [Key initialization]
+        1. Data loader configuration (optimizing performance parameters)
+        2. Control group expression value (used for direction loss calculation)
+        3. dict_filter: Non-zero gene index corresponding to each perturbation (affects the training loss calculation range)
+        """
+        self.weight_bias_track = weight_bias_track
+        self.is_master = is_master
+        self.local_rank = local_rank
+        self.world_size = world_size
+        self.proj_name = proj_name
+        self.exp_name = exp_name
+
+        # Weights & Biases Experiment Tracking
+        if self.weight_bias_track:
+            import wandb
+            wandb.init(project=proj_name, name=exp_name, entity = "deanghoss")
+            self.wandb = wandb
+        else:
+            self.wandb = None
+
+        self.device = f'cuda:{local_rank}'
+        torch.cuda.set_device(device)
+        self.config = None
+
+        # ============================================================
+        # [Performance optimization] Data loader configuration
+        # pin_memory: Accelerate CPU->GPU transfer
+        # prefetch_factor: prefetch multiple batches
+        # persistent_workers: Keep worker processes alive to avoid restart overhead
+        # ============================================================
+        num_workers = min(os.cpu_count() // world_size, 8)  # Up to 8 workers per GPU
+
+        self.dataloader = {}
+        for dl_key, dl in pert_data.dataloader.items():
+            dataset = dl.dataset
+            if 'train' in dl_key:
+                sampler = DistributedSampler(dataset, shuffle=True, drop_last=True)
+                self.dataloader[dl_key] = DataLoader(
+                    dataset,
+                    batch_size=train_bs,
+                    sampler=sampler,
+                    num_workers=num_workers,
+                    pin_memory=True,
+                    prefetch_factor=2,
+                    persistent_workers=True if num_workers > 0 else False,
+                    drop_last=True  # Make sure all batch sizes are consistent
+                )
+            elif 'val' in dl_key:
+                sampler = SequentialDistributedSampler(dataset, batch_size=test_bs, world_size=self.world_size)
+                self.dataloader[dl_key] = DataLoader(
+                    dataset,
+                    batch_size=test_bs,
+                    sampler=sampler,
+                    num_workers=min(num_workers, 4),
+                    pin_memory=True,
+                    prefetch_factor=2,
+                    persistent_workers=True if num_workers > 0 else False
+                )
+            else:
+                self.dataloader[dl_key] = dl
+
+        # Save data reference
+        self.adata = pert_data.adata
+        self.node_map = pert_data.node_map
+        self.data_path = pert_data.data_path
+        self.dataset_name = pert_data.dataset_name
+        self.split = pert_data.split
+        self.seed = pert_data.seed
+        self.train_gene_set_size = pert_data.train_gene_set_size
+        self.set2conditions = pert_data.set2conditions
+        self.subgroup = pert_data.subgroup
+        self.gene_list = list(pert_data.gene_names)
+        self.num_genes = len(self.gene_list)
+
+        try:
+            self.real_gene_list = list(pert_data.real_gene_names)
+        except AttributeError:
+            self.real_gene_list = list(pert_data.gene_names)
+        self.real_num_genes = len(self.real_gene_list)
+
+        if self.num_genes != self.real_num_genes:
+            print_sys("The gene list is augmented, larger than real gene list!")
+
+        # ============================================================
+        # [Key]Control group expression value
+        # Used to calculate direction loss (whether the predicted up/down direction is correct)
+        # ============================================================
+        self.ctrl_expression = torch.tensor(np.mean(self.adata.X[self.adata.obs.condition == 'ctrl'], axis = 0)).reshape(-1,).to(self.device)
+
+        # ============================================================
+        # [Key] dict_filter: Non-zero gene index of each perturbation
+        # The training loss is only calculated on these genes to avoid dropout noise
+        # This is an important setting that affects the results!
+        # ============================================================
+        pert_full_id2pert = dict(self.adata.obs[['condition_name', 'condition']].values)
+        self.dict_filter = {pert_full_id2pert[i]: j for i,j in
+                            self.adata.uns['non_zeros_gene_idx'].items() if
+                            i in pert_full_id2pert}
+        self.ctrl_adata = self.adata[self.adata.obs['condition'] == 'ctrl']
+
+        # Loss record
+        self.train_losses = []
+        self.train_overall_mses = []
+        self.val_overall_mses = []
+        self.train_de_mses = []
+        self.val_de_mses = []
+        
+    def tunable_parameters(self):
+        """
+        Returns a list of tunable parameters and their descriptions
+
+        [Description of key parameters]
+        - hidden_size: embedding dimension, affecting model capacity
+        - num_go_gnn_layers: Number of GO-GNN layers, which affects the propagation range of disturbance information
+        - coexpress_threshold: Threshold for co-expression graph construction, affecting the sparseness of the graph
+        - direction_lambda: direction loss weight, affecting the accuracy of upward/downward prediction
+        """
+        return {'hidden_size': 'hidden dimension, default 64',
+                'num_go_gnn_layers': 'number of GNN layers for GO graph, default 1',
+                'num_gene_gnn_layers': 'number of GNN layers for co-expression gene graph, default 1',
+                'decoder_hidden_size': 'hidden dimension for gene-specific decoder, default 16',
+                'num_similar_genes_go_graph': 'number of maximum similar K genes in the GO graph, default 20',
+                'num_similar_genes_co_express_graph': 'number of maximum similar K genes in the co expression graph, default 20',
+                'coexpress_threshold': 'pearson correlation threshold when constructing coexpression graph, default 0.4',
+                'uncertainty': 'whether or not to turn on uncertainty mode, default False',
+                'uncertainty_reg': 'regularization term to balance uncertainty loss and prediction loss, default 1',
+                'direction_lambda': 'regularization term to balance direction loss and prediction loss, default 1'
+               }
+
+    def model_initialize(self, hidden_size = 64,
+                         num_go_gnn_layers = 1,
+                         num_gene_gnn_layers = 1,
+                         decoder_hidden_size = 16,
+                         num_similar_genes_go_graph = 20,
+                         num_similar_genes_co_express_graph = 20,
+                         coexpress_threshold = 0.4,
+                         uncertainty = False,
+                         uncertainty_reg = 1,
+                         direction_lambda = 1e-1,
+                         G_go = None,
+                         G_go_weight = None,
+                         G_coexpress = None,
+                         G_coexpress_weight = None,
+                         no_perturb = False,
+                         cell_fitness_pred = False,
+                         go_path = None,
+                         model_type=None,
+                         bin_set=None,
+                         load_path=None,
+                         finetune_method=None,
+                         accumulation_steps=1,
+                         mode='v1',
+                         input_genes_ens_ids = None,
+                         scfm_genes_ens_ids = None,
+                         scfm_hyper_params_path = None,
+                         scfm_ckpt_path = 'None',
+                         scfm_class = None,
+                         model_class = None,
+                         key_enc = None,
+                         **kwargs
+                        ):
+        """
+        Model initialization
+
+        [Key steps]
+        1. Build a configuration dictionary (config is a single source of truth)
+        2. Construct co-expression graph and GO similarity graph (if not provided)
+        3. Instantiate the model
+
+        [Key parameters that affect the results]
+        - finetune_method: 'frozen' freezes pre-training parameters, None means fine-tuning all parameters
+        - accumulation_steps: the number of gradient accumulation steps, affecting the equivalent batch size
+        - coexpress_threshold/num_similar_genes_*: affects the structure of gene relationship graph
+
+        [Construction of genetic relationship diagram]
+        - G_coexpress: Data-driven co-expression relationship
+        - G_go: Functional similarity relationship based on GO annotation
+        These two graphs are important structural priors of the model and affect the way disturbance information is propagated.
+        """
+        # config is a single source of truth, and subsequent train/test/save all rely on the keys here.
+        self.config = {'proj_name': self.proj_name,
+                       'exp_name': self.exp_name,
+                        'hidden_size': hidden_size,
+                       'num_go_gnn_layers' : num_go_gnn_layers,
+                       'num_gene_gnn_layers' : num_gene_gnn_layers,
+                       'decoder_hidden_size' : decoder_hidden_size,
+                       'num_similar_genes_go_graph' : num_similar_genes_go_graph,
+                       'num_similar_genes_co_express_graph' : num_similar_genes_co_express_graph,
+                       'coexpress_threshold': coexpress_threshold,
+                       'uncertainty' : uncertainty,
+                       'uncertainty_reg' : uncertainty_reg,
+                       'direction_lambda' : direction_lambda,
+                       'G_go': G_go,
+                       'G_go_weight': G_go_weight,
+                       'G_coexpress': G_coexpress,
+                       'G_coexpress_weight': G_coexpress_weight,
+                       'device': self.device,
+                       'num_genes': self.num_genes,
+                       'no_perturb': no_perturb,
+                       'cell_fitness_pred': cell_fitness_pred,
+                        'model_type': model_type,
+                        'bin_set': bin_set,
+                        'load_path': load_path,
+                        'finetune_method': finetune_method,
+                        'accumulation_steps': accumulation_steps,
+                        'mode':mode,
+                        'input_genes_ens_ids': input_genes_ens_ids,
+                        'scfm_genes_ens_ids': scfm_genes_ens_ids,
+                        "scfm_hyper_params_path": scfm_hyper_params_path,
+                        "scfm_ckpt": scfm_ckpt_path,
+                        "scfm_class": scfm_class,
+                        "key_enc": key_enc
+                      }
+        print('Use accumulation steps:',accumulation_steps)
+        print('Use mode:',mode)
+        self.config.update(kwargs)
+
+        if self.wandb:
+            self.wandb.config.update(self.config)
+
+        # ============================================================
+        # [Key] Constructing a co-expression graph
+        # Only use training set data to avoid information leakage
+        # The graph structure relies on split/seed, which affects the reproducibility of results.
+        # ============================================================
+        if self.config['G_coexpress'] is None:
+            edge_list = get_similarity_network(network_type = 'co-express', adata = self.adata, threshold = coexpress_threshold, k = num_similar_genes_co_express_graph, gene_list = self.gene_list, data_path = self.data_path, data_name = self.dataset_name, split = self.split, seed = self.seed, train_gene_set_size = self.train_gene_set_size, set2conditions = self.set2conditions)
+            sim_network = GeneSimNetwork(edge_list, self.gene_list, node_map = self.node_map)
+            self.config['G_coexpress'] = sim_network.edge_index
+            self.config['G_coexpress_weight'] = sim_network.edge_weight
+
+        # ============================================================
+        # [Key] Constructing GO similarity graph
+        # This is an important source of priors for model structure
+        # Help models understand gene functional relationships and guide the propagation of perturbation effects
+        # ============================================================
+        if self.config['G_go'] is None:
+            edge_list = get_similarity_network(network_type = 'go', adata = self.adata, threshold = coexpress_threshold, k = num_similar_genes_go_graph, gene_list = self.gene_list, data_path = self.data_path, data_name = self.dataset_name, split = self.split, seed = self.seed, train_gene_set_size = self.train_gene_set_size, set2conditions = self.set2conditions)
+            print_sys(f"The go_edge_list shape: {edge_list.shape}")
+            sim_network = GeneSimNetwork(edge_list, self.gene_list, node_map = self.node_map)
+            self.config['G_go'] = sim_network.edge_index
+            self.config['G_go_weight'] = sim_network.edge_weight
+
+        # Instantiate model
+        self.model = model_class(self.config).to(self.device)
+        
+    def load_pretrained(self, path):
+        """
+        Load pretrained model
+
+        [Function]
+        Load configuration and model weights from saved directory
+
+        [Processing DDP prefix]
+        If the key of state_dict starts with "module.", it means it is a DDP model.
+        Need to remove prefix to load correctly
+        """
+        with open(os.path.join(path, 'config.pkl'), 'rb') as f:
+            config = pickle.load(f)
+
+        del config['device'], config['num_genes'], config['num_perts']
+        self.model_initialize(**config)
+        self.config = config
+
+        state_dict = torch.load(os.path.join(path, 'model.pt'), map_location = torch.device('cpu'))
+        # Handle the "module." prefix of DDP models
+        if next(iter(state_dict))[:7] == 'module.':
+            from collections import OrderedDict
+            new_state_dict = OrderedDict()
+            for k, v in state_dict.items():
+                name = k[7:] # remove `module.`
+                new_state_dict[name] = v
+            state_dict = new_state_dict
+
+        self.model.load_state_dict(state_dict)
+        self.model = self.model.to(self.device)
+        self.best_model = self.model
+
+    def save_model(self, path):
+        """
+        Save model
+
+        [Save content]
+        1. config.pkl: configuration dictionary (hyperparameters, paths, etc.)
+        2. model.pt: model weights
+
+        [Notice]
+        What is saved is the module of the DDP model, not the DDP wrapper itself.
+        """
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+        if self.config is None:
+            raise ValueError('No model is initialized...')
+
+        with open(os.path.join(path, 'config.pkl'), 'wb') as f:
+            pickle.dump(self.config, f)
+
+        torch.save(self.model.module.state_dict(), os.path.join(path, 'model.pt'))
+
+    def predict(self, pert_list):
+        """
+        Predicting transcriptome expression for a given combination of perturbations
+
+        [Function]
+        Given a list of perturbed genes, return the predicted expression value after perturbation
+
+        [parameter]
+        - pert_list: perturbation list, each element is a gene list
+          For example: [['MAPK1'], ['KRAS'], ['MAPK1', 'KRAS']]
+
+        [return]
+        - Predicted gene expression dictionary
+        - If uncertainty is turned on, the uncertainty score will also be returned
+        """
+        self.ctrl_adata = self.adata[self.adata.obs['condition'] == 'ctrl']
+        for pert in pert_list:
+            for i in pert:
+                if i not in self.pert_list:
+                    raise ValueError(i+ " not in the perturbation graph. Please select from PertNet.gene_list!")
+
+        if self.config['uncertainty']:
+            results_logvar = {}
+
+        self.best_model = self.best_model.to(self.device)
+        self.best_model.eval()
+        results_pred = {}
+        results_logvar_sum = {}
+
+        from torch_geometric.data import DataLoader
+        for pert in pert_list:
+            try:
+                # If prediction already exists, skip inference
+                results_pred['_'.join(pert)] = self.saved_pred['_'.join(pert)]
+                if self.config['uncertainty']:
+                    results_logvar_sum['_'.join(pert)] = self.saved_logvar_sum['_'.join(pert)]
+                continue
+            except:
+                pass
+
+            # Create cytogram data for prediction
+            cg = create_cell_graph_dataset_for_prediction(pert, self.ctrl_adata, self.pert_list, self.device)
+            loader = DataLoader(cg, 5, shuffle = False)
+
+            predall=[]
+            for step, batch in enumerate(loader):
+                batch.to(self.device)
+                with torch.no_grad():
+                    if self.config['uncertainty']:
+                        p, unc = self.best_model(batch)
+                        results_logvar['_'.join(pert)] = np.mean(unc.detach().cpu().numpy(), axis = 0)
+                        results_logvar_sum['_'.join(pert)] = np.exp(-np.mean(results_logvar['_'.join(pert)]))
+                    else:
+                        p = self.best_model(batch)
+                    predall.append(p.detach().cpu().numpy())
+            # Averaging predictions from multiple control group samples
+            preadall = np.concatenate(predall,axis=0)
+            results_pred['_'.join(pert)] = np.mean(preadall, axis = 0)
+
+        self.saved_pred.update(results_pred)
+
+        if self.config['uncertainty']:
+            self.saved_logvar_sum.update(results_logvar_sum)
+            return results_pred, results_logvar_sum
+        else:
+            return results_pred
+
+    def GI_predict(self, combo, GI_genes_file='./genes_with_hi_mean.npy'):
+        """
+        Predict the genetic interaction effect of gene pairs
+
+        [Function]
+        Given two genes, predict the interactive effects of their combined perturbations
+
+        [Genetic interaction analysis]
+        - Analyze whether two genes act synergistically, antagonistically, or independently
+        - Calculate linear combination coefficients, distance correlation coefficients and other indicators
+
+        [parameter]
+        - combo: gene pair, such as ('MAPK1', 'KRAS')
+        - GI_genes_file: Gene list file used for GI analysis
+        """
+        try:
+            # If there is already a prediction, use it directly
+            pred = {}
+            pred[combo[0]] = self.saved_pred[combo[0]]
+            pred[combo[1]] = self.saved_pred[combo[1]]
+            pred['_'.join(combo)] = self.saved_pred['_'.join(combo)]
+        except:
+            if self.config['uncertainty']:
+                pred = self.predict([[combo[0]], [combo[1]], combo])[0]
+            else:
+                pred = self.predict([[combo[0]], [combo[1]], combo])
+
+        # Calculate change relative to control group
+        mean_control = get_mean_control(self.adata).values
+        pred = {p:pred[p]-mean_control for p in pred}
+
+        if GI_genes_file is not None:
+            GI_genes_idx = get_GI_genes_idx(self.adata, GI_genes_file)
+        else:
+            GI_genes_idx = np.arange(len(self.adata.var.gene_name.values))
+
+        pred = {p:pred[p][GI_genes_idx] for p in pred}
+        return get_GI_params(pred, combo)
+
+    def plot_perturbation(self, query, save_file = None):
+        """
+        Visualizing disturbance prediction results
+
+        [Function]
+        Draw a comparison chart between predicted values and actual values to show the prediction effect of top DE genes
+
+        [Chart description]
+        - Box plot: distribution of true values
+        - Red dot: predicted value
+        - Green dashed line: control group baseline
+        """
+        import seaborn as sns
+        import numpy as np
+        import matplotlib.pyplot as plt
+
+        sns.set_theme(style="ticks", rc={"axes.facecolor": (0, 0, 0, 0)}, font_scale=1.5)
+
+        adata = self.adata
+        gene2idx = self.node_map
+        cond2name = dict(adata.obs[['condition', 'condition_name']].values)
+        gene_raw2id = dict(zip(adata.var.index.values, adata.var.gene_name.values))
+
+        de_idx = [gene2idx[gene_raw2id[i]] for i in adata.uns['top_non_dropout_de_20'][cond2name[query]]]
+        genes = [gene_raw2id[i] for i in adata.uns['top_non_dropout_de_20'][cond2name[query]]]
+        truth = adata[adata.obs.condition == query].X.toarray()[:, de_idx]
+        pred = self.predict([query.split('+')])['_'.join(query.split('+'))][de_idx]
+        ctrl_means = adata[adata.obs['condition'] == 'ctrl'].to_df().mean()[de_idx].values
+
+        pred = pred - ctrl_means
+        truth = truth - ctrl_means
+
+        plt.figure(figsize=[16.5,4.5])
+        plt.title(query)
+        plt.boxplot(truth, showfliers=False,
+                    medianprops = dict(linewidth=0))
+
+        for i in range(pred.shape[0]):
+            _ = plt.scatter(i+1, pred[i], color='red')
+
+        plt.axhline(0, linestyle="dashed", color = 'green')
+
+        ax = plt.gca()
+        ax.xaxis.set_ticklabels(genes, rotation = 90)
+
+        plt.ylabel("Change in Gene Expression over Control",labelpad=10)
+        plt.tick_params(axis='x', which='major', pad=5)
+        plt.tick_params(axis='y', which='major', pad=5)
+        sns.despine()
+
+        if save_file:
+            plt.savefig(save_file, bbox_inches='tight')
+        plt.show()
+    
+    
+    def train(self, 
+              epochs = 20, 
+              result_dir='./results',
+              lr = 1e-3,
+              weight_decay = 5e-4, 
+              valid_every = 1
+             ):
+        # Simulate larger batches through gradient accumulation. This value affects optimization stability and equivalent learning rate.
+        GRADIENT_ACCUMULATION = self.config['accumulation_steps']
+        test = (epochs == 1)
+        train_loader = self.dataloader['train_loader']
+        val_loader = self.dataloader['val_loader']
+            
+        self.model = self.model
+
+        # Key switch: Freezing the pre-trained encoder reduces overfitting, but reduces adaptation capabilities.
+        if self.config['finetune_method'] == 'frozen':
+            for name, p in self.model.named_parameters():
+                if "singlecell_model" in name:
+                    p.requires_grad = False
+        else:
+            for name, p in self.model.named_parameters():
+                p.requires_grad = True
+
+        # ============================================================
+        # Optimization: DDP configuration optimization
+        # ============================================================
+        self.model = DDP(self.model,
+                        device_ids=[self.local_rank],
+                        output_device=self.local_rank,
+                        find_unused_parameters=True,  # True is required because a few parameters are not used (such as gene_emb when pretrained=True)
+                        broadcast_buffers=False,  # Optimization: skip buffer broadcast (such as BN running stats)
+                        gradient_as_bucket_view=True)  # Optimization: reduce memory copy
+
+        optimizer = optim.Adam(self.model.parameters(), lr=lr, weight_decay = weight_decay)
+        scheduler = StepLR(optimizer, step_size=1, gamma=0.5)
+
+        min_val = np.inf
+        dist.barrier()
+        if self.is_master:
+            print_sys('Start Training...')
+        for epoch in range(epochs):
+            train_loader.sampler.set_epoch(epoch)
+            self.model.train()
+            dist.barrier()
+            for step, batch in enumerate(train_loader):
+                if test and (step+1 > 2 * GRADIENT_ACCUMULATION):
+                    break
+                batch.to(self.device)
+                y = batch.y
+                if (step+1) % GRADIENT_ACCUMULATION != 0:
+                    # no_sync: Avoid doing all-reduce for each micro-batch to save communication.
+                    with self.model.no_sync():
+                        pred = self.model(batch)
+                        loss = loss_fct(pred, y, batch.pert,
+                                    ctrl = self.ctrl_expression, 
+                                    dict_filter = self.dict_filter,
+                                    direction_lambda = self.config['direction_lambda'])
+                    
+                        loss.backward()
+                elif (step+1) % GRADIENT_ACCUMULATION == 0:
+                    pred = self.model(batch)
+                    loss = loss_fct(pred, y, batch.pert,
+                                ctrl = self.ctrl_expression, 
+                                dict_filter = self.dict_filter,
+                                direction_lambda = self.config['direction_lambda'])
+                
+                    loss.backward()
+                    # Gradient clipping affects training stability; if it is too small, it will inhibit learning, and if it is too large, the gradient may explode.
+                    nn.utils.clip_grad_value_(self.model.parameters(), clip_value=1.0)
+                    optimizer.step()
+                    optimizer.zero_grad()
+                
+                if (step + 1) % 50 == 0:
+                    if self.is_master:
+                        if self.wandb:
+                            self.wandb.log({'training_loss': loss.item()})
+                        log = "Epoch {} Step {} Train Loss: {:.4f}" 
+                        print_sys(log.format(epoch + 1, step + 1, loss.item()))
+            dist.barrier()
+            scheduler.step()
+            if (epoch + 1) % valid_every == 0:
+                # Evaluate model performance on train and val set
+                with torch.no_grad():
+                    val_res = evaluate(val_loader, self.model, self.config['uncertainty'], self.device)
+                val_metrics, _ = compute_metrics(val_res)
+                del val_res
+                val_metrics_reduced = {}
+                for m in ['mse', 'mse_de', 'pearson', 'pearson_de']:
+                    val_metrics_reduced[m] = get_reduced(val_metrics[m], self.local_rank, 0, self.world_size)
+                val_metrics = val_metrics_reduced
+                # Print epoch performance
+                log = "Epoch {}: Validation Overall MSE: {:.4f}. Validation Top 20 DE MSE: {:.4f}."
+
+                if self.is_master:
+                    print_sys(log.format(epoch + 1, val_metrics['mse'], val_metrics['mse_de']))
+                    if self.wandb:
+                        metrics = ['mse', 'pearson']
+                        for m in metrics:
+                            self.wandb.log({'val_'+m: val_metrics[m],'val_de_'+m: val_metrics[m + '_de']})
+                    # Use `mse_de` as early stop/select optimal criterion: pay more attention to DE gene prediction quality.
+                    if val_metrics['mse_de'] < min_val:
+                        min_val = val_metrics['mse_de']
+                        print_sys("Best epoch:{} mse_de:{}!".format(epoch + 1, min_val))
+                        self.save_model(result_dir)
+            
+        dist.barrier()
+        if self.is_master:
+            print_sys("Done!")
+            if 'test_loader' not in self.dataloader:
+                print_sys('Done! No test dataloader detected.')
+                return
+            # Force the best checkpoint to be loaded before testing and avoid using the parameters of the last epoch.
+            checkpoint = torch.load(os.path.join(result_dir, 'model.pt'), map_location=self.device)
+            self.model.module.load_state_dict(checkpoint)
+            self.model.eval()
+
+            if self.config.get('record_pred', False):
+                self.model.module.args['start_record'] = True
+                
+            # Model testing
+            test_loader = self.dataloader['test_loader']
+            if self.is_master:
+                print_sys("Start Testing...")
+            with torch.no_grad():
+                test_res = evaluate(test_loader, self.model.module, self.config['uncertainty'], self.device)
+                '''
+                if not test:
+                    test_res = evaluate(test_loader, self.model.module, self.config['uncertainty'], self.device)
+                else:
+                    quick_model = debug_model()
+                    test_res = evaluate(test_loader, quick_model, self.config['uncertainty'], self.device)'''
+            test_metrics, test_pert_res = compute_metrics(test_res)
+            
+            '''
+            test_metrics_reduced = {}
+            for m in test_metrics.keys():
+                test_metrics_reduced[m] = get_reduced(test_metrics[m], self.local_rank, 0, self.world_size)
+            test_metrics = test_metrics_reduced'''
+
+            log = "Best performing model: Test Top 20 DE MSE: {:.4f}"
+            if self.is_master:
+                print_sys(log.format(test_metrics['mse_de']))
+                if self.wandb:
+                    metrics = ['mse', 'pearson']
+                    for m in metrics:
+                        self.wandb.log({'test_' + m: test_metrics[m],
+                                'test_de_'+m: test_metrics[m + '_de']                     
+                                })
+                    
+            out = deeper_analysis(self.adata, test_res)
+            out_non_dropout = non_dropout_analysis(self.adata, test_res) 
+
+            if self.split == 'simulation':
+                print_sys("Start doing subgroup analysis for simulation split...")
+                subgroup = self.subgroup
+                subgroup_analysis = {}
+                for name in subgroup['test_subgroup'].keys():
+                    subgroup_analysis[name] = {}
+                    for m in list(list(test_pert_res.values())[0].keys()):
+                        subgroup_analysis[name][m] = []
+
+                for name, pert_list in subgroup['test_subgroup'].items():
+                    for pert in pert_list:
+                        for m, res in test_pert_res[pert].items():
+                            subgroup_analysis[name][m].append(res)
+
+                for name, result in subgroup_analysis.items():
+                    for m in result.keys():
+                        subgroup_analysis[name][m] = np.mean(subgroup_analysis[name][m])
+                        if self.wandb:
+                            self.wandb.log({'test_' + name + '_' + m: subgroup_analysis[name][m]})
+
+                        print_sys('test_' + name + '_' + m + ': ' + str(subgroup_analysis[name][m]))
+
+                ## deeper analysis
+                subgroup_analysis = {}
+                for name in subgroup['test_subgroup'].keys():
+                    subgroup_analysis[name] = {}
+
+                for name, pert_list in subgroup['test_subgroup'].items():
+                    for pert in pert_list:
+                        for m in out[pert].keys():
+                            if m in subgroup_analysis[name].keys():
+                                subgroup_analysis[name][m].append(out[pert][m])
+                            else:
+                                subgroup_analysis[name][m] = [out[pert][m]]
+
+                        for m in out_non_dropout[pert].keys():
+                            if m in subgroup_analysis[name].keys():
+                                subgroup_analysis[name][m].append(out_non_dropout[pert][m])
+                            else:
+                                subgroup_analysis[name][m] = [out_non_dropout[pert][m]]
+
+                for name, result in subgroup_analysis.items():
+                    for m in result.keys():
+                        subgroup_analysis[name][m] = np.mean(subgroup_analysis[name][m])
+                        if self.wandb:
+                            self.wandb.log({'test_' + name + '_' + m: subgroup_analysis[name][m]})
+
+                        print_sys('test_' + name + '_' + m + ': ' + str(subgroup_analysis[name][m]))
+            print_sys('Done!')
+
+
+            
+

BioMod/gene_perturbation/inference_local.py

@@ -0,0 +1,1034 @@
+"""
+================================================================================
+Model evaluation module (inference_local.py)
+================================================================================
+
+[Module function]
+Provides evaluation and analysis functions for genetic perturbation models, including:
+
+1. evaluate(): main evaluation function, calculates MSE, Pearson correlation coefficient and other indicators
+2. compute_metrics(): Calculate the complete dictionary of evaluation metrics
+3. deeper_analysis(): In-depth analysis of prediction results
+4. non_dropout_analysis(): analyze on non-dropout genes
+
+[Description of key indicators]
+- MSE: mean square error of all genes
+- MSE_DE: MSE of Top 20 differentially expressed genes (paying more attention to biological significance)
+- Pearson: Pearson correlation coefficient of all genes
+- Pearson_DE: Pearson correlation coefficient of Top 20 DE genes
+
+[Evaluation Notes]
+- real_num_genes: only evaluate on real genes, excluding extended/completed genes
+- de_idx: The differentially expressed gene index of each sample may be different
+================================================================================
+"""
+
+import torch
+import numpy as np
+import scanpy as sc
+
+from sklearn.metrics import r2_score
+from scipy.stats import pearsonr, spearmanr
+from sklearn.metrics import mean_squared_error as mse
+from sklearn.metrics import mean_absolute_error as mae
+
+def evaluate(loader, model, uncertainty, device, real_num_genes = None):
+    """
+    main evaluation function
+
+    [Function]
+    Runs the model on the given data loader, collecting predicted and true values.
+
+    [Key steps]
+    1. Traverse all batches of the data loader
+    2. Run model inference (without calculating gradients)
+    3. Collect predicted values, true values, and DE gene predicted values
+
+    [Parameter description]
+    - real_num_genes: only evaluate the first N real genes to avoid expanding gene interference indicators
+    """
+
+    model.eval()
+    model.to(device)
+    pert_cat = []
+    pred = []
+    truth = []
+    pred_de = []
+    truth_de = []
+    results = {}
+    logvar = []
+    
+    for itr, batch in enumerate(loader):
+
+        batch.to(device)
+        pert_cat.extend(batch.pert)
+
+        with torch.no_grad():
+            if uncertainty:
+                p, unc = model(batch)
+                logvar.extend(unc.cpu())
+            else:
+                p = model(batch)
+            # Only evaluate in the real gene dimension to avoid the interference of completed/expanded genes on indicators.
+            p = p[:, :real_num_genes]
+            t = batch.y[:, :real_num_genes]
+            pred.extend(p.cpu())
+            truth.extend(t.cpu())
+            
+            # The DE indicator is only counted on the de_idx corresponding to the sample and is one of the core calibers of downstream reporting.
+            for itr, de_idx in enumerate(batch.de_idx):
+                pred_de.append(p[itr, de_idx])
+                truth_de.append(t[itr, de_idx])
+
+    # all genes
+    results['pert_cat'] = np.array(pert_cat)
+    pred = torch.stack(pred)
+    truth = torch.stack(truth)
+    results['pred']= pred.detach().cpu().numpy()
+    results['truth']= truth.detach().cpu().numpy()
+
+    pred_de = torch.stack(pred_de)
+    truth_de = torch.stack(truth_de)
+    results['pred_de']= pred_de.detach().cpu().numpy()
+    results['truth_de']= truth_de.detach().cpu().numpy()
+    
+    if uncertainty:
+        results['logvar'] = torch.stack(logvar).detach().cpu().numpy()
+    
+    return results
+
+#This evaluation is just as test. IT HAS INFORMATION LEAKAGE!
+def evaluate_with_nonzero_mask(loader, model, uncertainty, device):
+    """
+    Run model in inference mode using a given data loader
+    """
+
+    model.eval()
+    model.to(device)
+    pert_cat = []
+    pred = []
+    truth = []
+    pred_de = []
+    truth_de = []
+    results = {}
+    logvar = []
+    
+    for itr, batch in enumerate(loader):
+
+        batch.to(device)
+        pert_cat.extend(batch.pert)
+
+        with torch.no_grad():
+            if uncertainty:
+                p, unc = model(batch)
+                logvar.extend(unc.cpu())
+            else:
+                p = model(batch)
+            t = batch.y
+            p = p * (t != 0)
+            pred.extend(p.cpu())
+            truth.extend(t.cpu())
+            
+            # Differentially expressed genes
+            for itr, de_idx in enumerate(batch.de_idx):
+                pred_de.append(p[itr, de_idx])
+                truth_de.append(t[itr, de_idx])
+
+    # all genes
+    results['pert_cat'] = np.array(pert_cat)
+    pred = torch.stack(pred)
+    truth = torch.stack(truth)
+    results['pred']= pred.detach().cpu().numpy()
+    results['truth']= truth.detach().cpu().numpy()
+
+    pred_de = torch.stack(pred_de)
+    truth_de = torch.stack(truth_de)
+    results['pred_de']= pred_de.detach().cpu().numpy()
+    results['truth_de']= truth_de.detach().cpu().numpy()
+    
+    if uncertainty:
+        results['logvar'] = torch.stack(logvar).detach().cpu().numpy()
+    
+    return results
+
+
+def non_zero_evaluate(loader, model, uncertainty, device, 
+                        delta_list, penalty = {
+                            "false_neg": 1,
+                            "false_pos": 1 }):
+    """
+    Run model in inference mode using a given data loader
+    """
+
+    model.eval()
+    model.to(device)
+    pert_cat = []
+    pred = []
+    truth = []
+    pred_de = []
+    truth_de = []
+    results = {}
+    logvar = []
+    
+    for itr, batch in enumerate(loader):
+
+        batch.to(device)
+        pert_cat.extend(batch.pert)
+
+        with torch.no_grad():
+            if uncertainty:
+                p, unc = model(batch)
+                logvar.extend(unc.cpu())
+            else:
+                p = model(batch)
+            t = batch.y
+            pred.extend(p.cpu())
+            truth.extend(t.cpu())
+            
+            # Differentially expressed genes
+            for itr, de_idx in enumerate(batch.de_idx):
+                pred_de.append(p[itr, de_idx])
+                truth_de.append(t[itr, de_idx])
+
+    # all genes
+    results['pert_cat'] = np.array(pert_cat)
+    pred = torch.stack(pred)
+    truth = torch.stack(truth)
+    results['pred']= pred.detach().cpu().numpy()
+    results['truth']= truth.detach().cpu().numpy()
+
+    pred_de = torch.stack(pred_de)
+    truth_de = torch.stack(truth_de)
+    results['pred_de']= pred_de.detach().cpu().numpy()
+    results['truth_de']= truth_de.detach().cpu().numpy()
+    
+    if uncertainty:
+        results['logvar'] = torch.stack(logvar).detach().cpu().numpy()
+    
+    return results
+
+
+def compute_metrics(results, gene_idx=None):
+    """
+    Given results from a model run and the ground truth, compute metrics
+
+    """
+    metrics = {}
+    metrics_pert = {}
+
+    metric2fct = {
+           'mse': mse,
+           'pearson': pearsonr
+    }
+    
+    for m in metric2fct.keys():
+        metrics[m] = []
+        metrics[m + '_de'] = []
+
+    for pert in np.unique(results['pert_cat']):
+
+        metrics_pert[pert] = {}
+        p_idx = np.where(results['pert_cat'] == pert)[0]
+            
+        for m, fct in metric2fct.items():
+            # First average the cells of the same perturbation, and then calculate the indicator of the perturbation.
+            if m == 'pearson':
+                val = fct(results['pred'][p_idx].mean(0), results['truth'][p_idx].mean(0))[0]
+                if np.isnan(val):
+                    val = 0
+            else:
+                val = fct(results['pred'][p_idx].mean(0), results['truth'][p_idx].mean(0))
+
+            metrics_pert[pert][m] = val
+            metrics[m].append(metrics_pert[pert][m])
+
+       
+        if pert != 'ctrl':
+            
+            for m, fct in metric2fct.items():
+                if m == 'pearson':
+                    val = fct(results['pred_de'][p_idx].mean(0), results['truth_de'][p_idx].mean(0))[0]
+                    if np.isnan(val):
+                        val = 0
+                else:
+                    val = fct(results['pred_de'][p_idx].mean(0), results['truth_de'][p_idx].mean(0))
+                    
+                metrics_pert[pert][m + '_de'] = val
+                metrics[m + '_de'].append(metrics_pert[pert][m + '_de'])
+
+        else:
+            for m, fct in metric2fct.items():
+                metrics_pert[pert][m + '_de'] = 0
+    
+    for m in metric2fct.keys():
+        # The result is a perturbation-level macro average (each perturbation is equally weighted).
+        metrics[m] = np.mean(metrics[m])
+        metrics[m + '_de'] = np.mean(metrics[m + '_de'])
+    
+    return metrics, metrics_pert
+
+
+def compute_synergy_loss(results, mean_control, high_umi_idx, subtype = 'POTENTIATION'):
+    pred_res = get_test_set_results_seen2(results, subtype)
+    all_perts = np.unique(results['pert_cat'])
+    linear_params = get_linear_params(pred_res, high_umi_idx,
+                                            mean_control,all_perts)
+    synergy_loss = np.sum([np.abs(linear_params[k]['pred']['mag']
+                               - linear_params[k]['truth']['mag']) for k in
+                                linear_params])
+    
+    mag = np.sum([np.abs(linear_params[k]['pred']['mag']) for k in
+                                linear_params])
+    return synergy_loss, mag
+
+
+def non_zero_analysis(adata, test_res):
+    metric2fct = {
+           'pearson': pearsonr,
+           'mse': mse
+    }
+
+    pert_metric = {}
+    
+    ## in silico modeling and upperbounding
+    pert2pert_full_id = dict(adata.obs[['condition', 'condition_name']].values)
+    geneid2name = dict(zip(adata.var.index.values, adata.var['gene_name']))
+    geneid2idx = dict(zip(adata.var.index.values, range(len(adata.var.index.values))))
+
+    # calculate mean expression for each condition
+    unique_conditions = adata.obs.condition.unique()
+    conditions2index = {}
+    for i in unique_conditions:
+        conditions2index[i] = np.where(adata.obs.condition == i)[0]
+
+    condition2mean_expression = {}
+    for i, j in conditions2index.items():
+        condition2mean_expression[i] = np.mean(adata.X[j], axis = 0)
+    pert_list = np.array(list(condition2mean_expression.keys()))
+    mean_expression = np.array(list(condition2mean_expression.values())).reshape(len(adata.obs.condition.unique()), adata.X.toarray().shape[1])
+    ctrl = mean_expression[np.where(pert_list == 'ctrl')[0]]
+    
+    gene_list = adata.var['gene_name'].values
+
+    for pert in np.unique(test_res['pert_cat']):
+        pert_metric[pert] = {}
+        
+        pert_idx = np.where(test_res['pert_cat'] == pert)[0]    
+        de_idx = [geneid2idx[i] for i in adata.uns['top_non_zero_de_20'][pert2pert_full_id[pert]]]
+
+        direc_change = np.abs(np.sign(test_res['pred'][pert_idx].mean(0)[de_idx] - ctrl[0][de_idx]) - np.sign(test_res['truth'][pert_idx].mean(0)[de_idx] - ctrl[0][de_idx]))            
+        frac_correct_direction = len(np.where(direc_change == 0)[0])/len(de_idx)
+        pert_metric[pert]['frac_correct_direction_top20_non_zero'] = frac_correct_direction
+        
+        frac_direction_opposite = len(np.where(direc_change == 2)[0])/len(de_idx)
+        pert_metric[pert]['frac_opposite_direction_top20_non_zero'] = frac_direction_opposite
+        
+        frac_direction_opposite = len(np.where(direc_change == 1)[0])/len(de_idx)
+        pert_metric[pert]['frac_0/1_direction_top20_non_zero'] = frac_direction_opposite
+        
+        mean = np.mean(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        std = np.std(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        min_ = np.min(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        max_ = np.max(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        q25 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.25, axis = 0)
+        q75 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.75, axis = 0)
+        q55 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.55, axis = 0)
+        q45 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.45, axis = 0)
+        q40 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.4, axis = 0)
+        q60 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.6, axis = 0)
+        
+        zero_des = np.intersect1d(np.where(min_ == 0)[0], np.where(max_ == 0)[0])
+        nonzero_des = np.setdiff1d(list(range(20)), zero_des)
+        
+        if len(nonzero_des) == 0:
+            pass
+            # pert that all de genes are 0...
+        else:            
+            pred_mean = np.mean(test_res['pred'][pert_idx][:, de_idx], axis = 0).reshape(-1,)
+            true_mean = np.mean(test_res['truth'][pert_idx][:, de_idx], axis = 0).reshape(-1,)
+           
+            in_range = (pred_mean[nonzero_des] >= min_[nonzero_des]) & (pred_mean[nonzero_des] <= max_[nonzero_des])
+            frac_in_range = sum(in_range)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_non_zero'] = frac_in_range
+
+            in_range_5 = (pred_mean[nonzero_des] >= q45[nonzero_des]) & (pred_mean[nonzero_des] <= q55[nonzero_des])
+            frac_in_range_45_55 = sum(in_range_5)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_45_55_non_zero'] = frac_in_range_45_55
+
+            in_range_10 = (pred_mean[nonzero_des] >= q40[nonzero_des]) & (pred_mean[nonzero_des] <= q60[nonzero_des])
+            frac_in_range_40_60 = sum(in_range_10)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_40_60_non_zero'] = frac_in_range_40_60
+
+            in_range_25 = (pred_mean[nonzero_des] >= q25[nonzero_des]) & (pred_mean[nonzero_des] <= q75[nonzero_des])
+            frac_in_range_25_75 = sum(in_range_25)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_25_75_non_zero'] = frac_in_range_25_75
+
+            zero_idx = np.where(std > 0)[0]
+            sigma = (np.abs(pred_mean[zero_idx] - mean[zero_idx]))/(std[zero_idx])
+            pert_metric[pert]['mean_sigma_non_zero'] = np.mean(sigma)
+            pert_metric[pert]['std_sigma_non_zero'] = np.std(sigma)
+            pert_metric[pert]['frac_sigma_below_1_non_zero'] = 1 - len(np.where(sigma > 1)[0])/len(zero_idx)
+            pert_metric[pert]['frac_sigma_below_2_non_zero'] = 1 - len(np.where(sigma > 2)[0])/len(zero_idx)
+        
+        p_idx = np.where(test_res['pert_cat'] == pert)[0]
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top20_de_non_zero'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx], test_res['truth'][p_idx].mean(0)[de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top20_de_non_zero'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])
+                pert_metric[pert][m + '_top20_de_non_zero'] = val
+                
+    return pert_metric
+
+def non_dropout_analysis(adata, test_res):
+    metric2fct = {
+           'pearson': pearsonr,
+           'mse': mse
+    }
+
+    pert_metric = {}
+    
+    ## in silico modeling and upperbounding
+    pert2pert_full_id = dict(adata.obs[['condition', 'condition_name']].values)
+    geneid2name = dict(zip(adata.var.index.values, adata.var['gene_name']))
+    geneid2idx = dict(zip(adata.var.index.values, range(len(adata.var.index.values))))
+
+    # calculate mean expression for each condition
+    unique_conditions = adata.obs.condition.unique()
+    conditions2index = {}
+    for i in unique_conditions:
+        conditions2index[i] = np.where(adata.obs.condition == i)[0]
+
+    condition2mean_expression = {}
+    for i, j in conditions2index.items():
+        condition2mean_expression[i] = np.mean(adata.X[j], axis = 0)
+    pert_list = np.array(list(condition2mean_expression.keys()))
+    mean_expression = np.array(list(condition2mean_expression.values())).reshape(len(adata.obs.condition.unique()), adata.X.toarray().shape[1])
+    ctrl = mean_expression[np.where(pert_list == 'ctrl')[0]]
+    
+    gene_list = adata.var['gene_name'].values
+
+    for pert in np.unique(test_res['pert_cat']):
+        pert_metric[pert] = {}
+        
+        pert_idx = np.where(test_res['pert_cat'] == pert)[0]    
+        de_idx = [geneid2idx[i] for i in adata.uns['top_non_dropout_de_20'][pert2pert_full_id[pert]]]
+        non_zero_idx = adata.uns['non_zeros_gene_idx'][pert2pert_full_id[pert]]
+        non_dropout_gene_idx = adata.uns['non_dropout_gene_idx'][pert2pert_full_id[pert]]
+             
+        direc_change = np.abs(np.sign(test_res['pred'][pert_idx].mean(0)[de_idx] - ctrl[0][de_idx]) - np.sign(test_res['truth'][pert_idx].mean(0)[de_idx] - ctrl[0][de_idx]))            
+        frac_correct_direction = len(np.where(direc_change == 0)[0])/len(de_idx)
+        pert_metric[pert]['frac_correct_direction_top20_non_dropout'] = frac_correct_direction
+        
+        frac_direction_opposite = len(np.where(direc_change == 2)[0])/len(de_idx)
+        pert_metric[pert]['frac_opposite_direction_top20_non_dropout'] = frac_direction_opposite
+        
+        frac_direction_opposite = len(np.where(direc_change == 1)[0])/len(de_idx)
+        pert_metric[pert]['frac_0/1_direction_top20_non_dropout'] = frac_direction_opposite
+        
+        direc_change = np.abs(np.sign(test_res['pred'][pert_idx].mean(0)[non_zero_idx] - ctrl[0][non_zero_idx]) - np.sign(test_res['truth'][pert_idx].mean(0)[non_zero_idx] - ctrl[0][non_zero_idx]))            
+        frac_correct_direction = len(np.where(direc_change == 0)[0])/len(non_zero_idx)
+        pert_metric[pert]['frac_correct_direction_non_zero'] = frac_correct_direction
+
+        frac_direction_opposite = len(np.where(direc_change == 2)[0])/len(non_zero_idx)
+        pert_metric[pert]['frac_opposite_direction_non_zero'] = frac_direction_opposite
+        
+        frac_direction_opposite = len(np.where(direc_change == 1)[0])/len(non_zero_idx)
+        pert_metric[pert]['frac_0/1_direction_non_zero'] = frac_direction_opposite
+        
+        direc_change = np.abs(np.sign(test_res['pred'][pert_idx].mean(0)[non_dropout_gene_idx] - ctrl[0][non_dropout_gene_idx]) - np.sign(test_res['truth'][pert_idx].mean(0)[non_dropout_gene_idx] - ctrl[0][non_dropout_gene_idx]))            
+        frac_correct_direction = len(np.where(direc_change == 0)[0])/len(non_dropout_gene_idx)
+        pert_metric[pert]['frac_correct_direction_non_dropout'] = frac_correct_direction
+        
+        frac_direction_opposite = len(np.where(direc_change == 2)[0])/len(non_dropout_gene_idx)
+        pert_metric[pert]['frac_opposite_direction_non_dropout'] = frac_direction_opposite
+        
+        frac_direction_opposite = len(np.where(direc_change == 1)[0])/len(non_dropout_gene_idx)
+        pert_metric[pert]['frac_0/1_direction_non_dropout'] = frac_direction_opposite
+        
+        mean = np.mean(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        std = np.std(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        min_ = np.min(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        max_ = np.max(test_res['truth'][pert_idx][:, de_idx], axis = 0)
+        q25 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.25, axis = 0)
+        q75 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.75, axis = 0)
+        q55 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.55, axis = 0)
+        q45 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.45, axis = 0)
+        q40 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.4, axis = 0)
+        q60 = np.quantile(test_res['truth'][pert_idx][:, de_idx], 0.6, axis = 0)
+        
+        zero_des = np.intersect1d(np.where(min_ == 0)[0], np.where(max_ == 0)[0])
+        nonzero_des = np.setdiff1d(list(range(20)), zero_des)
+        
+        if len(nonzero_des) == 0:
+            pass
+            # pert that all de genes are 0...
+        else:            
+            pred_mean = np.mean(test_res['pred'][pert_idx][:, de_idx], axis = 0).reshape(-1,)
+            true_mean = np.mean(test_res['truth'][pert_idx][:, de_idx], axis = 0).reshape(-1,)
+           
+            in_range = (pred_mean[nonzero_des] >= min_[nonzero_des]) & (pred_mean[nonzero_des] <= max_[nonzero_des])
+            frac_in_range = sum(in_range)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_non_dropout'] = frac_in_range
+
+            in_range_5 = (pred_mean[nonzero_des] >= q45[nonzero_des]) & (pred_mean[nonzero_des] <= q55[nonzero_des])
+            frac_in_range_45_55 = sum(in_range_5)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_45_55_non_dropout'] = frac_in_range_45_55
+
+            in_range_10 = (pred_mean[nonzero_des] >= q40[nonzero_des]) & (pred_mean[nonzero_des] <= q60[nonzero_des])
+            frac_in_range_40_60 = sum(in_range_10)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_40_60_non_dropout'] = frac_in_range_40_60
+
+            in_range_25 = (pred_mean[nonzero_des] >= q25[nonzero_des]) & (pred_mean[nonzero_des] <= q75[nonzero_des])
+            frac_in_range_25_75 = sum(in_range_25)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_25_75_non_dropout'] = frac_in_range_25_75
+
+            zero_idx = np.where(std > 0)[0]
+            sigma = (np.abs(pred_mean[zero_idx] - mean[zero_idx]))/(std[zero_idx])
+            pert_metric[pert]['mean_sigma_non_dropout'] = np.mean(sigma)
+            pert_metric[pert]['std_sigma_non_dropout'] = np.std(sigma)
+            pert_metric[pert]['frac_sigma_below_1_non_dropout'] = 1 - len(np.where(sigma > 1)[0])/len(zero_idx)
+            pert_metric[pert]['frac_sigma_below_2_non_dropout'] = 1 - len(np.where(sigma > 2)[0])/len(zero_idx)
+        
+        p_idx = np.where(test_res['pert_cat'] == pert)[0]
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top20_de_non_dropout'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx], test_res['truth'][p_idx].mean(0)[de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top20_de_non_dropout'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])
+                pert_metric[pert][m + '_top20_de_non_dropout'] = val
+                
+    return pert_metric
+    
+def deeper_analysis(adata, test_res, de_column_prefix = 'rank_genes_groups_cov', most_variable_genes = None):
+    
+    metric2fct = {
+           'pearson': pearsonr,
+           'mse': mse
+    }
+
+    pert_metric = {}
+
+    ## in silico modeling and upperbounding
+    pert2pert_full_id = dict(adata.obs[['condition', 'condition_name']].values)
+    geneid2name = dict(zip(adata.var.index.values, adata.var['gene_name']))
+    geneid2idx = dict(zip(adata.var.index.values, range(len(adata.var.index.values))))
+
+    # calculate mean expression for each condition
+    unique_conditions = adata.obs.condition.unique()
+    conditions2index = {}
+    for i in unique_conditions:
+        conditions2index[i] = np.where(adata.obs.condition == i)[0]
+
+    condition2mean_expression = {}
+    for i, j in conditions2index.items():
+        condition2mean_expression[i] = np.mean(adata.X[j], axis = 0)
+    pert_list = np.array(list(condition2mean_expression.keys()))
+    mean_expression = np.array(list(condition2mean_expression.values())).reshape(len(adata.obs.condition.unique()), adata.X.toarray().shape[1])
+    ctrl = mean_expression[np.where(pert_list == 'ctrl')[0]]
+    
+    if most_variable_genes is None:
+        most_variable_genes = np.argsort(np.std(mean_expression, axis = 0))[-200:]
+        
+    gene_list = adata.var['gene_name'].values
+
+    for pert in np.unique(test_res['pert_cat']):
+        pert_metric[pert] = {}
+        de_idx = [geneid2idx[i] for i in adata.uns['rank_genes_groups_cov_all'][pert2pert_full_id[pert]][:20]]
+        de_idx_200 = [geneid2idx[i] for i in adata.uns['rank_genes_groups_cov_all'][pert2pert_full_id[pert]][:200]]
+        de_idx_100 = [geneid2idx[i] for i in adata.uns['rank_genes_groups_cov_all'][pert2pert_full_id[pert]][:100]]
+        de_idx_50 = [geneid2idx[i] for i in adata.uns['rank_genes_groups_cov_all'][pert2pert_full_id[pert]][:50]]
+
+        pert_idx = np.where(test_res['pert_cat'] == pert)[0]    
+        pred_mean = np.mean(test_res['pred_de'][pert_idx], axis = 0).reshape(-1,)
+        true_mean = np.mean(test_res['truth_de'][pert_idx], axis = 0).reshape(-1,)
+        
+        direc_change = np.abs(np.sign(test_res['pred'][pert_idx].mean(0) - ctrl[0]) - np.sign(test_res['truth'][pert_idx].mean(0) - ctrl[0]))            
+        frac_correct_direction = len(np.where(direc_change == 0)[0])/len(geneid2name)
+        pert_metric[pert]['frac_correct_direction_all'] = frac_correct_direction
+
+        de_idx_map = {20: de_idx,
+                      50: de_idx_50,
+                      100: de_idx_100,
+                      200: de_idx_200
+                     }
+        
+        for val in [20, 50, 100, 200]:
+            
+            direc_change = np.abs(np.sign(test_res['pred'][pert_idx].mean(0)[de_idx_map[val]] - ctrl[0][de_idx_map[val]]) - np.sign(test_res['truth'][pert_idx].mean(0)[de_idx_map[val]] - ctrl[0][de_idx_map[val]]))            
+            frac_correct_direction = len(np.where(direc_change == 0)[0])/val
+            pert_metric[pert]['frac_correct_direction_' + str(val)] = frac_correct_direction
+
+        mean = np.mean(test_res['truth_de'][pert_idx], axis = 0)
+        std = np.std(test_res['truth_de'][pert_idx], axis = 0)
+        min_ = np.min(test_res['truth_de'][pert_idx], axis = 0)
+        max_ = np.max(test_res['truth_de'][pert_idx], axis = 0)
+        q25 = np.quantile(test_res['truth_de'][pert_idx], 0.25, axis = 0)
+        q75 = np.quantile(test_res['truth_de'][pert_idx], 0.75, axis = 0)
+        q55 = np.quantile(test_res['truth_de'][pert_idx], 0.55, axis = 0)
+        q45 = np.quantile(test_res['truth_de'][pert_idx], 0.45, axis = 0)
+        q40 = np.quantile(test_res['truth_de'][pert_idx], 0.4, axis = 0)
+        q60 = np.quantile(test_res['truth_de'][pert_idx], 0.6, axis = 0)
+
+        zero_des = np.intersect1d(np.where(min_ == 0)[0], np.where(max_ == 0)[0])
+        nonzero_des = np.setdiff1d(list(range(20)), zero_des)
+        if len(nonzero_des) == 0:
+            pass
+            # pert that all de genes are 0...
+        else:            
+            
+            direc_change = np.abs(np.sign(pred_mean[nonzero_des] - ctrl[0][de_idx][nonzero_des]) - np.sign(true_mean[nonzero_des] - ctrl[0][de_idx][nonzero_des]))            
+            frac_correct_direction = len(np.where(direc_change == 0)[0])/len(nonzero_des)
+            pert_metric[pert]['frac_correct_direction_20_nonzero'] = frac_correct_direction
+            
+            in_range = (pred_mean[nonzero_des] >= min_[nonzero_des]) & (pred_mean[nonzero_des] <= max_[nonzero_des])
+            frac_in_range = sum(in_range)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range'] = frac_in_range
+
+            in_range_5 = (pred_mean[nonzero_des] >= q45[nonzero_des]) & (pred_mean[nonzero_des] <= q55[nonzero_des])
+            frac_in_range_45_55 = sum(in_range_5)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_45_55'] = frac_in_range_45_55
+
+            in_range_10 = (pred_mean[nonzero_des] >= q40[nonzero_des]) & (pred_mean[nonzero_des] <= q60[nonzero_des])
+            frac_in_range_40_60 = sum(in_range_10)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_40_60'] = frac_in_range_40_60
+
+            in_range_25 = (pred_mean[nonzero_des] >= q25[nonzero_des]) & (pred_mean[nonzero_des] <= q75[nonzero_des])
+            frac_in_range_25_75 = sum(in_range_25)/len(nonzero_des)
+            pert_metric[pert]['frac_in_range_25_75'] = frac_in_range_25_75
+
+            zero_idx = np.where(std > 0)[0]
+            sigma = (np.abs(pred_mean[zero_idx] - mean[zero_idx]))/(std[zero_idx])
+            pert_metric[pert]['mean_sigma'] = np.mean(sigma)
+            pert_metric[pert]['std_sigma'] = np.std(sigma)
+            pert_metric[pert]['frac_sigma_below_1'] = 1 - len(np.where(sigma > 1)[0])/len(zero_idx)
+            pert_metric[pert]['frac_sigma_below_2'] = 1 - len(np.where(sigma > 2)[0])/len(zero_idx)
+
+        ## correlation on delta
+        p_idx = np.where(test_res['pert_cat'] == pert)[0]
+
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+                val = fct(test_res['pred'][p_idx].mean(0)- ctrl[0], test_res['truth'][p_idx].mean(0)-ctrl[0])[0]
+                if np.isnan(val):
+                    val = 0
+
+                pert_metric[pert][m + '_delta'] = val
+                
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+
+                pert_metric[pert][m + '_delta_de'] = val
+
+        ## up fold changes > 10, 
+        pert_mean = np.mean(test_res['truth'][p_idx], axis = 0).reshape(-1,)
+
+        fold_change = pert_mean/ctrl
+        fold_change[np.isnan(fold_change)] = 0
+        fold_change[np.isinf(fold_change)] = 0
+        ## this is to remove the ones that are super low and the fold change becomes unmeaningful
+        fold_change[0][np.where(pert_mean < 0.5)[0]] = 0
+
+        o =  np.where(fold_change[0] > 0)[0]
+
+        pred_fc = test_res['pred'][p_idx].mean(0)[o]
+        true_fc = test_res['truth'][p_idx].mean(0)[o]
+        ctrl_fc = ctrl[0][o]
+
+        if len(o) > 0:
+            pert_metric[pert]['fold_change_gap_all'] = np.mean(np.abs(pred_fc/ctrl_fc - true_fc/ctrl_fc))
+
+
+        o = np.intersect1d(np.where(fold_change[0] <0.333)[0], np.where(fold_change[0] > 0)[0])
+
+        pred_fc = test_res['pred'][p_idx].mean(0)[o]
+        true_fc = test_res['truth'][p_idx].mean(0)[o]
+        ctrl_fc = ctrl[0][o]
+
+        if len(o) > 0:
+            pert_metric[pert]['fold_change_gap_downreg_0.33'] = np.mean(np.abs(pred_fc/ctrl_fc - true_fc/ctrl_fc))
+
+
+        o = np.intersect1d(np.where(fold_change[0] <0.1)[0], np.where(fold_change[0] > 0)[0])
+
+        pred_fc = test_res['pred'][p_idx].mean(0)[o]
+        true_fc = test_res['truth'][p_idx].mean(0)[o]
+        ctrl_fc = ctrl[0][o]
+
+        if len(o) > 0:
+            pert_metric[pert]['fold_change_gap_downreg_0.1'] = np.mean(np.abs(pred_fc/ctrl_fc - true_fc/ctrl_fc))
+
+        o = np.where(fold_change[0] > 3)[0]
+
+        pred_fc = test_res['pred'][p_idx].mean(0)[o]
+        true_fc = test_res['truth'][p_idx].mean(0)[o]
+        ctrl_fc = ctrl[0][o]
+
+        if len(o) > 0:
+            pert_metric[pert]['fold_change_gap_upreg_3'] = np.mean(np.abs(pred_fc/ctrl_fc - true_fc/ctrl_fc))
+
+        o = np.where(fold_change[0] > 10)[0]
+
+        pred_fc = test_res['pred'][p_idx].mean(0)[o]
+        true_fc = test_res['truth'][p_idx].mean(0)[o]
+        ctrl_fc = ctrl[0][o]
+
+        if len(o) > 0:
+            pert_metric[pert]['fold_change_gap_upreg_10'] = np.mean(np.abs(pred_fc/ctrl_fc - true_fc/ctrl_fc))
+
+        ## most variable genes
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+                val = fct(test_res['pred'][p_idx].mean(0)[most_variable_genes] - ctrl[0][most_variable_genes], test_res['truth'][p_idx].mean(0)[most_variable_genes]-ctrl[0][most_variable_genes])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top200_hvg'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[most_variable_genes], test_res['truth'][p_idx].mean(0)[most_variable_genes])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top200_hvg'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[most_variable_genes], test_res['truth'][p_idx].mean(0)[most_variable_genes])
+                pert_metric[pert][m + '_top200_hvg'] = val
+
+
+        ## top 20/50/100/200 DEs
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top20_de'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx], test_res['truth'][p_idx].mean(0)[de_idx])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top20_de'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx] - ctrl[0][de_idx], test_res['truth'][p_idx].mean(0)[de_idx]-ctrl[0][de_idx])
+                pert_metric[pert][m + '_top20_de'] = val
+
+        
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_200] - ctrl[0][de_idx_200], test_res['truth'][p_idx].mean(0)[de_idx_200]-ctrl[0][de_idx_200])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top200_de'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_200], test_res['truth'][p_idx].mean(0)[de_idx_200])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top200_de'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_200] - ctrl[0][de_idx_200], test_res['truth'][p_idx].mean(0)[de_idx_200]-ctrl[0][de_idx_200])
+                pert_metric[pert][m + '_top200_de'] = val
+
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_100] - ctrl[0][de_idx_100], test_res['truth'][p_idx].mean(0)[de_idx_100]-ctrl[0][de_idx_100])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top100_de'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_100], test_res['truth'][p_idx].mean(0)[de_idx_100])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top100_de'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_100] - ctrl[0][de_idx_100], test_res['truth'][p_idx].mean(0)[de_idx_100]-ctrl[0][de_idx_100])
+                pert_metric[pert][m + '_top100_de'] = val
+
+        for m, fct in metric2fct.items():
+            if m != 'mse':
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_50] - ctrl[0][de_idx_50], test_res['truth'][p_idx].mean(0)[de_idx_50]-ctrl[0][de_idx_50])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_delta_top50_de'] = val
+
+
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_50], test_res['truth'][p_idx].mean(0)[de_idx_50])[0]
+                if np.isnan(val):
+                    val = 0
+                pert_metric[pert][m + '_top50_de'] = val
+            else:
+                val = fct(test_res['pred'][p_idx].mean(0)[de_idx_50] - ctrl[0][de_idx_50], test_res['truth'][p_idx].mean(0)[de_idx_50]-ctrl[0][de_idx_50])
+                pert_metric[pert][m + '_top50_de'] = val
+
+
+
+    return pert_metric
+
+
+def GI_subgroup(pert_metric):
+    GI_type2Score = {}
+    test_pert_list = list(pert_metric.keys())
+    for GI_type, gi_list in GIs.items():
+        intersect = np.intersect1d(gi_list, test_pert_list)
+        if len(intersect) != 0:
+            GI_type2Score[GI_type] = {}
+
+            for m in list(list(pert_metric.values())[0].keys()):
+                GI_type2Score[GI_type][m] = np.mean([pert_metric[i][m] for i in intersect if m in pert_metric[i]])
+                
+    return GI_type2Score
+
+def node_specific_batch_out(models, batch):
+    # Returns output for all node specific models as a matrix of dimension batch_size x nodes
+    outs = []
+    for idx in range(len(models)):
+        outs.append(models[idx](batch).detach().cpu().numpy()[:,idx])
+    return np.vstack(outs).T
+
+# Run prediction over all batches
+def batch_predict(loader, loaded_models, args):
+    # Prediction for node specific GNNs
+    preds = []
+    print("Loader size: ", len(loader))
+    for itr, batch in enumerate(loader):
+        print(itr)
+        batch = batch.to(args['device'])
+        preds.append(node_specific_batch_out(loaded_models, batch))
+
+    preds = np.vstack(preds)
+    return preds
+
+def get_high_umi_idx(gene_list):
+    # Genes used for linear model fitting
+    try:
+        high_umi = np.load('../genes_with_hi_mean.npy', allow_pickle=True)
+    except:
+        high_umi = np.load('./genes_with_hi_mean.npy', allow_pickle=True)
+    high_umi_idx = np.where([g in high_umi for g in gene_list])[0]
+    return high_umi_idx
+
+def get_mean_ctrl(adata):
+    return adata[adata.obs['condition'] == 'ctrl'].to_df().mean().reset_index(
+        drop=True)
+
+def get_single_name(g, all_perts):
+    name = g+'+ctrl'
+    if name in all_perts:
+        return name
+    else:
+        return 'ctrl+'+g
+
+
+def get_test_set_results_seen2(res, sel_GI_type):
+    # Get relevant test set results
+    test_pert_cats = [p for p in np.unique(res['pert_cat']) if
+                      p in GIs[sel_GI_type] or 'ctrl' in p]
+    pred_idx = np.where([t in test_pert_cats for t in res['pert_cat']])
+    out = {}
+    for key in res:
+        out[key] = res[key][pred_idx]
+    return out
+
+## Synergy loss calculation functions
+def get_all_vectors(all_res, mean_control, double,
+                    single1, single2, high_umi_idx):
+    # Pred
+    pred_df = pd.DataFrame(all_res['pred'])
+    pred_df['condition'] = all_res['pert_cat']
+    subset_df = pred_df[pred_df['condition'] == double].iloc[:, :-1]
+    delta_double_pred = subset_df.mean(0) - mean_control
+    single_df_1_pred = pred_df[pred_df['condition'] == single1].iloc[:, :-1]
+    single_df_2_pred = pred_df[pred_df['condition'] == single2].iloc[:, :-1]
+
+    # True
+    truth_df = pd.DataFrame(all_res['truth'])
+    truth_df['condition'] = all_res['pert_cat']
+    subset_df = truth_df[truth_df['condition'] == double].iloc[:, :-1]
+    delta_double_truth = subset_df.mean(0) - mean_control
+    single_df_1_truth = truth_df[truth_df['condition'] == single1].iloc[:, :-1]
+    single_df_2_truth = truth_df[truth_df['condition'] == single2].iloc[:, :-1]
+
+    delta_single_truth_1 = single_df_1_truth.mean(0) - mean_control
+    delta_single_truth_2 = single_df_2_truth.mean(0) - mean_control
+    delta_single_pred_1 = single_df_1_pred.mean(0) - mean_control
+    delta_single_pred_2 = single_df_2_pred.mean(0) - mean_control
+
+    return {'single_pred_1': delta_single_pred_1.values[high_umi_idx],
+            'single_pred_2': delta_single_pred_2.values[high_umi_idx],
+            'double_pred': delta_double_pred.values[high_umi_idx],
+            'single_truth_1': delta_single_truth_1.values[high_umi_idx],
+            'single_truth_2': delta_single_truth_2.values[high_umi_idx],
+            'double_truth': delta_double_truth.values[high_umi_idx]}
+
+
+def get_coeffs_synergy(singles_expr, double_expr):
+    results = {}
+    results['ts'] = TheilSenRegressor(fit_intercept=False,
+                                      max_subpopulation=1e5,
+                                      max_iter=1000,
+                                      random_state=1000)
+    X = singles_expr
+    y = double_expr
+    try:
+        results['ts'].fit(X, y.ravel())
+    except:
+        print(X)
+        print(y)
+    results['c1'] = results['ts'].coef_[0]
+    results['c2'] = results['ts'].coef_[1]
+    results['mag'] = np.sqrt((results['c1'] ** 2 + results['c2'] ** 2))
+    return results
+
+
+def Fit(all_vectors, type_='pertnet'):
+    if type_ == 'pertnet':
+        singles_expr = np.array(
+            [all_vectors['single_pred_1'], all_vectors['single_pred_2']]).T
+        first_expr = np.array([all_vectors['single_pred_1']]).T
+        second_expr = np.array([all_vectors['single_pred_2']]).T
+        double_expr = np.array(all_vectors['double_pred']).T
+
+    elif type_ == 'truth':
+        singles_expr = np.array(
+            [all_vectors['single_truth_1'], all_vectors['single_truth_2']]).T
+        first_expr = np.array([all_vectors['single_truth_1']]).T
+        second_expr = np.array([all_vectors['single_truth_2']]).T
+        double_expr = np.array(all_vectors['double_truth']).T
+
+    return get_coeffs_synergy(singles_expr, double_expr)
+
+
+def get_linear_params(pred_res, high_umi_idx, mean_control, all_perts):
+    results = {}
+    for c in set(pred_res['pert_cat']):
+        if 'ctrl' in c:
+            continue
+        double = c
+        single1 = get_single_name(double.split('+')[0], all_perts)
+        single2 = get_single_name(double.split('+')[1], all_perts)
+        all_vectors = get_all_vectors(pred_res, mean_control, double,
+                                      single1, single2, high_umi_idx)
+
+        pertnet_results = Fit(all_vectors, type_='pertnet')
+        truth_results = Fit(all_vectors, type_='truth')
+
+        results[c] = {
+            'truth': truth_results,
+            'pred': pertnet_results}
+
+    return results
+
+
+# Read in model for each gene
+
+GIs = {
+    'NEOMORPHIC': ['CBL+TGFBR2',
+                  'KLF1+TGFBR2',
+                  'MAP2K6+SPI1',
+                  'SAMD1+TGFBR2',
+                  'TGFBR2+C19orf26',
+                  'TGFBR2+ETS2',
+                  'CBL+UBASH3A',
+                  'CEBPE+KLF1',
+                  'DUSP9+MAPK1',
+                  'FOSB+PTPN12',
+                  'PLK4+STIL',
+                  'PTPN12+OSR2',
+                  'ZC3HAV1+CEBPE'],
+    'ADDITIVE': ['BPGM+SAMD1',
+                'CEBPB+MAPK1',
+                'CEBPB+OSR2',
+                'DUSP9+PRTG',
+                'FOSB+OSR2',
+                'IRF1+SET',
+                'MAP2K3+ELMSAN1',
+                'MAP2K6+ELMSAN1',
+                'POU3F2+FOXL2',
+                'RHOXF2BB+SET',
+                'SAMD1+PTPN12',
+                'SAMD1+UBASH3B',
+                'SAMD1+ZBTB1',
+                'SGK1+TBX2',
+                'TBX3+TBX2',
+                'ZBTB10+SNAI1'],
+    'EPISTASIS': ['AHR+KLF1',
+                 'MAPK1+TGFBR2',
+                 'TGFBR2+IGDCC3',
+                 'TGFBR2+PRTG',
+                 'UBASH3B+OSR2',
+                 'DUSP9+ETS2',
+                 'KLF1+CEBPA',
+                 'MAP2K6+IKZF3',
+                 'ZC3HAV1+CEBPA'],
+    'REDUNDANT': ['CDKN1C+CDKN1A',
+                 'MAP2K3+MAP2K6',
+                 'CEBPB+CEBPA',
+                 'CEBPE+CEBPA',
+                 'CEBPE+SPI1',
+                 'ETS2+MAPK1',
+                 'FOSB+CEBPE',
+                 'FOXA3+FOXA1'],
+    'POTENTIATION': ['CNN1+UBASH3A',
+                    'ETS2+MAP7D1',
+                    'FEV+CBFA2T3',
+                    'FEV+ISL2',
+                    'FEV+MAP7D1',
+                    'PTPN12+UBASH3A'],
+    'SYNERGY_SIMILAR_PHENO':['CBL+CNN1',
+                            'CBL+PTPN12',
+                            'CBL+PTPN9',
+                            'CBL+UBASH3B',
+                            'FOXA3+FOXL2',
+                            'FOXA3+HOXB9',
+                            'FOXL2+HOXB9',
+                            'UBASH3B+CNN1',
+                            'UBASH3B+PTPN12',
+                            'UBASH3B+PTPN9',
+                            'UBASH3B+ZBTB25'],
+    'SYNERGY_DISSIMILAR_PHENO': ['AHR+FEV',
+                                'DUSP9+SNAI1',
+                                'FOXA1+FOXF1',
+                                'FOXA1+FOXL2',
+                                'FOXA1+HOXB9',
+                                'FOXF1+FOXL2',
+                                'FOXF1+HOXB9',
+                                'FOXL2+MEIS1',
+                                'IGDCC3+ZBTB25',
+                                'POU3F2+CBFA2T3',
+                                'PTPN12+ZBTB25',
+                                'SNAI1+DLX2',
+                                'SNAI1+UBASH3B'],
+    'SUPPRESSOR': ['CEBPB+PTPN12',
+                  'CEBPE+CNN1',
+                  'CEBPE+PTPN12',
+                  'CNN1+MAPK1',
+                  'ETS2+CNN1',
+                  'ETS2+IGDCC3',
+                  'ETS2+PRTG',
+                  'FOSB+UBASH3B',
+                  'IGDCC3+MAPK1',
+                  'LYL1+CEBPB',
+                  'MAPK1+PRTG',
+                  'PTPN12+SNAI1']
+}

BioMod/gene_perturbation/model_module.py

@@ -0,0 +1,760 @@
+"""
+================================================================================
+Genetic perturbation task model definition (model_module.py)
+================================================================================
+
+[Model Architecture Overview]
+This document implements a deep learning model for genetic perturbation prediction. The core components include:
+
+1. ModuleEmbedding: Convert gene-level expression values into module-level tokens
+   - Supports multiple integration methods (multiply/add/gated_residual)
+   - Use attention pooling to aggregate genetic information within the module
+
+2. PerformerEncoderOnly_Module: Encoder-only Transformer model
+   - Use Performer (Linear Complexity Attention) to process large numbers of module tokens
+   - Support Masked Autoencoder (MAE) pre-training paradigm
+
+3. SCFM_Module: wrapper for pre-trained models
+   - Process the mapping of input genes to pre-trained model genes
+   - Recover gene-level embeddings from module encoding
+
+4. GEARS_Model_Module: Genetic perturbation prediction model
+   - Integrate pre-trained encoder + GNN + decoder
+   - Predict expression changes after gene perturbation
+
+[Key technical points]
+- Module-based gene representation: Group thousands of genes into hundreds of modules to reduce computational complexity
+- Gene embedding recovery: Recover gene-level features from module encoding through average pooling
+- GO graph neural network: using gene ontology knowledge to enhance perturbation representation
+
+================================================================================
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Sequential, Linear, ReLU
+import numpy as np
+import pickle
+from torch_geometric.nn import SGConv
+from typing import List, Tuple, Optional
+import sys
+import os
+from copy import deepcopy
+import scipy
+
+# ============================================================
+# Import Performer related components (does not depend on gears_reduce)
+# Performer is a linear complexity Transformer variant suitable for processing long sequences
+# ============================================================
+try:
+    from BioMod.models_performer_module import (
+        ModuleEmbedding,
+        PerformerEncoderOnly_Module,
+        sanitize_performer_module_config,
+    )
+except ModuleNotFoundError:
+    project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    if project_root not in sys.path:
+        sys.path.insert(0, project_root)
+    from BioMod.models_performer_module import (
+        ModuleEmbedding,
+        PerformerEncoderOnly_Module,
+        sanitize_performer_module_config,
+    )
+
+# ============================================================
+# Original GEARS related classes
+# ============================================================
+
+class MLP(torch.nn.Module):
+    """MLP network"""
+    def __init__(self, sizes, batch_norm=True, last_layer_act="linear"):
+        super(MLP, self).__init__()
+        layers = []
+        for s in range(len(sizes) - 1):
+            layers = layers + [
+                torch.nn.Linear(sizes[s], sizes[s + 1]),
+                torch.nn.BatchNorm1d(sizes[s + 1])
+                if batch_norm and s < len(sizes) - 1 else None,
+                torch.nn.ReLU()
+            ]
+
+        layers = [l for l in layers if l is not None][:-1]
+        self.activation = last_layer_act
+        self.network = torch.nn.Sequential(*layers)
+        self.relu = torch.nn.ReLU()
+    
+    def forward(self, x):
+        return self.network(x)
+
+
+def reindex_tensor_universal(tensor, index_positions, dim, filler='0', device='cpu'):
+    """
+    Reindex tensor along specified dimensions
+    """
+    index_tensor = torch.tensor(index_positions, device=device)
+    tensor_shape = list(tensor.shape)
+    tensor_shape[dim] = len(index_positions)
+    index_shape = [1] * len(tensor_shape)
+    index_shape[dim] = len(index_positions)
+    index_tensor = index_tensor.view(*index_shape).expand(*tensor_shape)
+
+    if filler == '0':
+        padder_shape = deepcopy(tensor_shape)
+        padder_shape[dim] = 1
+        padder = torch.zeros(padder_shape, dtype=tensor.dtype, device=device)
+    elif filler == 'mean':
+        padder = tensor.mean(dim=dim, keepdim=True)
+    else:
+        raise ValueError("filler should be 0 or mean!")
+    expanded_tensor = torch.cat((tensor, padder), dim=dim)
+    reindexed = torch.gather(expanded_tensor, dim, index_tensor)
+    return reindexed
+
+
+class SCFM_Module(torch.nn.Module):
+    """
+    ================================================================================
+    SCFM_Module: wrapper for pre-trained single cell base models
+
+    [Core functions]
+    Adapt the PerformerEncoderOnly_Module model to the genetic perturbation task to address the following key issues:
+
+    1. Gene mapping: The gene order of the input data may be different from that of the pre-trained model, and re-indexing is required.
+    2. Module-to-gene recovery: Recovering gene-level embeddings from module encodings
+    3. Dimension adjustment: adjust the encoder output to the dimensions required by downstream tasks
+
+    [Key Design Decisions]
+    - use_module_level=False (recommended): restore gene level embedding from module encoding
+      Advantages: Compatible with existing GEARS architecture, no need to modify downstream tasks
+    - use_module_level=True (experimental): use module level encoding directly
+      Downstream architecture needs to be redesigned
+
+    [Gene Embedding Recovery Strategy]
+    For each gene:
+    1. Find all modules containing the gene
+    2. Perform average pooling on the encoding of these modules
+    3. If the gene is not in any module, use the global average as fallback
+
+    [Performance Optimization]
+    - Precompute scatter index to avoid Python loops
+    - Vectorized aggregation using scatter_add
+    ================================================================================
+    """
+    def __init__(self, input_genes, scfm_genes,
+                 scfm_hyper_params, scfm_ckpt,
+                 module_sets=None,
+                 requires_grad=False,
+                 device='cpu',
+                 output_dim=200,
+                 use_module_level=False):
+        super(SCFM_Module, self).__init__()
+        self.device = device
+        self.use_module_level = use_module_level
+
+        # ============================================================
+        # [Gene Mapping] Establish index mapping between input genes and pre-trained model genes
+        # This is a critical step in dealing with inconsistencies in gene order across different datasets
+        # ============================================================
+        input_mapping = {idx: i for i, idx in enumerate(input_genes)}
+        self.input_gene_num = len(input_genes)
+
+        # ============================================================
+        # [Key] fusion_method processing
+        # Read from pkl to ensure consistency with pre-training
+        # ============================================================
+        fusion_method = scfm_hyper_params.get('fusion_method', 'multiply')
+        print(f"[SCFM_Module] Using fusion_method from hyper_params: {fusion_method}")
+
+        # Convert new indices to a list of index positions from the original tensor
+        scfm_genes_pad = scfm_genes + ['PAD']
+        self.scfm_seq_len = len(scfm_genes_pad)
+        scfm_mapping = {idx: i for i, idx in enumerate(scfm_genes_pad)}
+
+        # Create bidirectional index mapping
+        self.scfm_index_positions = [input_mapping.get(idx, self.input_gene_num) for idx in scfm_genes_pad]
+        self.input_index_positions = [scfm_mapping.get(idx, self.scfm_seq_len) for idx in input_genes]
+        
+        # Load module data (if not provided in hyper_params)
+        if module_sets is None:
+            if 'module_sets' in scfm_hyper_params and scfm_hyper_params['module_sets'] is not None:
+                module_sets = scfm_hyper_params['module_sets']
+            else:
+                raise ValueError("module_sets must be provided (via parameters or hyper_params)")
+
+        self.module_sets = module_sets
+        self.num_modules = len(module_sets)
+        self.scfm_seq_len = len(scfm_genes)
+
+        # ============================================================
+        # [Key] Create a gene-to-module mapping
+        # Each gene may appear in multiple modules for subsequent average pooling.
+        # This is the core data structure of module representation.
+        # ============================================================
+        self.gene_to_modules = {}  # {gene_idx: [module_idx1, module_idx2, ...]}
+        for module_idx, module_genes in enumerate(module_sets):
+            for gene_idx in module_genes:
+                if gene_idx < len(scfm_genes):  # Only consider genes in scfm_genes
+                    if gene_idx not in self.gene_to_modules:
+                        self.gene_to_modules[gene_idx] = []
+                    self.gene_to_modules[gene_idx].append(module_idx)
+
+        # ============================================================
+        # [Performance Optimization] Pre-built sparse index for efficient module->gene aggregation
+        # Use scatter_add to replace Python loop, greatly improving forward efficiency
+        # This is the standard optimization method in PyTorch for handling irregular aggregations
+        # ============================================================
+        import torch.sparse
+        max_modules_per_gene = max(len(v) for v in self.gene_to_modules.values()) if self.gene_to_modules else 1
+
+        # Construct a sparse index matrix: shape (num_genes, max_modules_per_gene)
+        gene_to_module_matrix = []
+        for gene_idx in range(self.scfm_seq_len):
+            if gene_idx in self.gene_to_modules:
+                modules = self.gene_to_modules[gene_idx]
+                padded = modules + [-1] * (max_modules_per_gene - len(modules))
+                gene_to_module_matrix.append(padded)
+            else:
+                gene_to_module_matrix.append([-1] * max_modules_per_gene)
+
+        # Registered as a buffer (will not be updated by the optimizer, but will be moved to the correct device with the model)
+        self.register_buffer('gene_to_module_matrix',
+                            torch.tensor(gene_to_module_matrix, dtype=torch.long))
+        gene_module_counts = torch.tensor([
+            len(self.gene_to_modules.get(i, [])) for i in range(self.scfm_seq_len)
+        ], dtype=torch.long)
+        self.register_buffer('gene_module_counts', gene_module_counts)
+
+        # ============================================================
+        # [Performance Optimization] Precompute scatter_add index
+        # Module->gene aggregation for fully vectorized, avoiding Python loops
+        # ============================================================
+        module_indices = []
+        gene_indices = []
+        for module_idx, module_genes in enumerate(module_sets):
+            for gene_idx in module_genes:
+                if gene_idx < self.scfm_seq_len:
+                    module_indices.append(module_idx)
+                    gene_indices.append(gene_idx)
+
+        self.register_buffer('scatter_module_idx', torch.tensor(module_indices, dtype=torch.long))
+        self.register_buffer('scatter_gene_idx', torch.tensor(gene_indices, dtype=torch.long))
+
+        # Create PerformerEncoderOnly_Module model
+        scfm_hyper_params['module_sets'] = module_sets
+        scfm_hyper_params['num_genes'] = len(scfm_genes)
+        scfm_hyper_params['max_seq_len'] = len(module_sets)
+
+        self.scfm = PerformerEncoderOnly_Module(**scfm_hyper_params)
+
+        # ============================================================
+        # [Key] Load pre-trained weights
+        # Filter out incompatible parameters (GO related, module_heads, etc.)
+        # strict=False allows partial parameter mismatch (used for fine-tuning mode)
+        # ============================================================
+        if scfm_ckpt != 'None':
+            ckpt = torch.load(scfm_ckpt, map_location=device)
+            model_state_dict = ckpt.get('model_state_dict', ckpt)
+            current_state_dict = self.scfm.state_dict()
+            filtered_state_dict = {}
+            for k, v in model_state_dict.items():
+                # Filter out GO related parameters (may not be available in different data sets)
+                if 'go_conv' in k or 'G_go' in k.lower():
+                    continue
+                # Filter out module_heads (not needed for fine-tuning)
+                if 'module_heads' in k:
+                    continue
+                # Filter out parameters that do not exist or have mismatched shapes in the current model
+                if k not in current_state_dict:
+                    continue
+                if current_state_dict[k].shape != v.shape:
+                    print(
+                        f"[SCFM_Module] skip shape-mismatched key: {k}, "
+                        f"ckpt={tuple(v.shape)} current={tuple(current_state_dict[k].shape)}"
+                    )
+                    continue
+                filtered_state_dict[k] = v
+            # Use strict=False because module_heads is intentionally removed
+            missing, unexpected = self.scfm.load_state_dict(filtered_state_dict, strict=False)
+            for param in self.scfm.parameters():
+                param.requires_grad = requires_grad
+        else:
+            for param in self.scfm.parameters():
+                param.requires_grad = True
+
+        # Output dimension adjustment (if needed)
+        self.output_dim = output_dim
+        if self.output_dim != 200:
+            try:
+                self.feedforward = nn.Linear(self.scfm.base_dim, self.output_dim)
+            except AttributeError:
+                self.feedforward = nn.Linear(scfm_hyper_params.get('base_dim', 200), self.output_dim)
+    
+    def forward(self, x):
+        """
+        Forward propagation: from gene expression values to gene-level embeddings
+
+        [Processing Process]
+        1. Re-index the input gene expression into the gene order of the pre-trained model
+        2. Get the module encoding through PerformerEncoderOnly_Module
+        3. Restore gene-level embedding from module encoding (or return module encoding directly)
+        4. Reindex gene embeddings back to input gene order
+
+        [Input and output]
+        Input: x (B, input_gene_num) - gene expression value
+        Output:
+            - use_module_level=False: output_reindex (B, input_gene_num, D) - gene level embedding
+            - use_module_level=True: module_encodings (B, num_modules, D) - module level encoding
+        """
+        x = x.view((-1, self.input_gene_num))  # B, M
+
+        # ============================================================
+        # [Step 1] Re-index: Map the input gene expression to the gene order of the pre-trained model
+        # Genes that are not present are filled with 0 (PAD position)
+        # ============================================================
+        x_scfm = reindex_tensor_universal(x, self.scfm_index_positions, dim=1, filler='0', device=self.device)  # B, N
+
+        # ============================================================
+        # [Step 2] Obtain the module code through the pre-trained model
+        # mask_ratio=0 means not to mask any modules (fine-tuning mode)
+        # Use mask_ratio>0 for MAE pre-training during pre-training.
+        # ============================================================
+        out = self.scfm(x_scfm, mask_ratio=0.0, return_encodings=True)
+
+        if 'encodings' in out:
+            module_encodings = out['encodings']  # Local legacy interface
+        elif 'module_encodings' in out:
+            module_encodings = out['module_encodings']  # Reuse pre-trained shared implementation interfaces
+        else:
+            raise ValueError(f"Expected one of ['encodings', 'module_encodings'] in output, got: {list(out.keys())}")
+
+        # ============================================================
+        # [Option] Directly return to module level encoding
+        # For experimental module-level downstream tasks
+        # ============================================================
+        if self.use_module_level:
+            if self.output_dim != 200:
+                B, num_modules, D = module_encodings.shape
+                module_encodings_flat = module_encodings.reshape(B * num_modules, D)
+                module_encodings_flat = self.feedforward(module_encodings_flat)
+                module_encodings = module_encodings_flat.reshape(B, num_modules, self.output_dim)
+            return module_encodings  # (B, num_modules, D)
+
+        # ============================================================
+        # [Step 3] Restore gene-level embedding from module encoding
+        # Core idea: The embedding of each gene = the average encoding of all modules containing the gene
+        # Vectorized aggregation using scatter_add to avoid Python loops
+        # ============================================================
+        B, num_modules, D = module_encodings.shape
+        G = self.scfm_seq_len
+
+        # Calculate global average (for genes without any modules)
+        global_mean = module_encodings.mean(dim=1)  # (B, D)
+
+        # Extract the corresponding module embeddings: (B, num_pairs, D)
+        module_embs = module_encodings[:, self.scatter_module_idx, :]  # (B, num_pairs, D)
+
+        # Aggregate to gene dimension using scatter_add
+        scfm_gene_embeddings = torch.zeros(B, G, D, device=module_encodings.device)
+        gene_idx_expanded = self.scatter_gene_idx.unsqueeze(0).unsqueeze(-1).expand(B, -1, D)
+        scfm_gene_embeddings.scatter_add_(1, gene_idx_expanded, module_embs)
+
+        # ============================================================
+        # [Key] Divide by the number of modules for each gene to get the average
+        # This is the core operation from module to gene aggregation
+        # clamp(min=1) avoids division by zero errors
+        # ============================================================
+        gene_counts = self.gene_module_counts.float().view(1, G, 1).clamp(min=1)
+        scfm_gene_embeddings = scfm_gene_embeddings / gene_counts
+
+        # For genes without any modules, use global average as fallback
+        no_module_mask = (self.gene_module_counts == 0).view(1, G, 1).float()
+        scfm_gene_embeddings = scfm_gene_embeddings * (1 - no_module_mask) + global_mean.unsqueeze(1) * no_module_mask
+
+        # Add embedding of PAD positions (for handling non-existent genes)
+        pad_emb = torch.zeros(B, 1, D, device=module_encodings.device)
+        output_encodings = torch.cat([scfm_gene_embeddings, pad_emb], dim=1)  # (B, G+1, D)
+
+        # ============================================================
+        # [Step 4] Reindex back to the input gene sequence
+        # Use mean as fallback to fill in non-existent genes
+        # ============================================================
+        output_reindex = reindex_tensor_universal(
+            output_encodings,
+            self.input_index_positions,
+            dim=1,
+            filler='mean',
+            device=self.device
+        )  # B, M, D
+
+        if self.output_dim != 200:
+            output_reindex = self.feedforward(output_reindex)
+
+        return output_reindex
+
+
+class GEARS_Model_Module(torch.nn.Module):
+    """
+    ================================================================================
+    GEARS_Model_Module: Genetic perturbation prediction model
+
+    [Core functions]
+    Predicting transcriptomic changes in cells following genetic perturbation. Enter a perturbation combination (such as knockout gene A),
+    Output the expression value prediction of each gene after perturbation.
+
+    [Model Architecture]
+    1. Pretrained encoder (SCFM_Module): Encode gene expression into embedding vectors
+    2. Perturbation embedding layer: generate perturbation features for perturbed genes
+    3. GO graph neural network (GNN): using gene ontology knowledge to enhance perturbation representation
+    4. Decoder: predict gene expression values after perturbation
+
+    [Key Design Decisions]
+    - Use GO graph neural network: gene ontology contains gene functional relationships, which can guide model understanding
+      How perturbation effects spread between genes
+    - Gene specific decoder: Each gene has independent decoding parameters to capture gene specificity
+    - Cross-gene MLP: Modeling interdependencies between genes
+
+    [Input and output]
+    Input: data (PyG Data object)
+        - x: (B*N, 2) feature matrix, [:, 0] is gene expression, [:, 1] is perturbation mark
+        - y: true expression value (used for training)
+    Output: out (B, N) - predicted gene expression value after perturbation
+    ================================================================================
+    """
+    def __init__(self, args):
+        super(GEARS_Model_Module, self).__init__()
+        self.args = args
+
+        # Forecast record related settings
+        if args.get('record_pred', False):
+            self.pred_dir = f"./result_process/preds/{self.args['exp_name']}"
+            if not os.path.exists(self.pred_dir):
+                os.makedirs(self.pred_dir)
+            self.pred_batch_idx = 0
+
+        # ============================================================
+        # [Basic parameters]
+        # hidden_size: embedding dimension, affecting model capacity
+        # num_genes: number of genes, determines the output dimension
+        # ============================================================
+        self.num_genes = args['num_genes']
+        hidden_size = args['hidden_size']
+        self.uncertainty = args['uncertainty']
+        self.num_layers = args['num_go_gnn_layers']  # GO GNN layers
+        self.indv_out_hidden_size = args['decoder_hidden_size']
+        self.num_layers_gene_pos = args['num_gene_gnn_layers']
+        self.pert_emb_lambda = 0.2
+
+        self.input_genes_list = args['input_genes_ens_ids']
+        self.scfm_genes_list = args['scfm_genes_ens_ids']
+
+        # ============================================================
+        # [Load pre-trained model configuration]
+        # hyper_params contains model architecture parameters, which must be consistent with pre-training
+        # ============================================================
+        try:
+            with open(args["scfm_hyper_params_path"], 'rb') as handle:
+                self.scfm_hyper_params = pickle.load(handle)
+        except ValueError:
+            import pickle5
+            with open(args["scfm_hyper_params_path"], 'rb') as handle:
+                self.scfm_hyper_params = pickle5.load(handle)
+        self.scfm_hyper_params.pop('module_embed_chunk_padded_gene_budget', None)
+        self.scfm_hyper_params = sanitize_performer_module_config(
+            self.scfm_hyper_params,
+            source=args["scfm_hyper_params_path"],
+            log_fn=print,
+        )
+
+        self.scfm_hyper_params['gene2vec_file'] = args['scfm_gene2vec_file']
+        self.scfm_ckpt = args["scfm_ckpt"]
+
+        # ============================================================
+        # [Module data loading]
+        # module_sets defines how genes are grouped into modules
+        # Can be loaded from hyper_params or external files
+        # ============================================================
+        module_sets = None
+        if 'module_sets' in self.scfm_hyper_params and self.scfm_hyper_params['module_sets'] is not None:
+            module_sets = self.scfm_hyper_params['module_sets']
+        elif args.get('module_data_path') or args.get('module_data_dir'):
+            from BioMod.module_dataset import load_module_data
+            module_sets = load_module_data(
+                module_data_path=args.get('module_data_path'),
+                module_data_dir=args.get('module_data_dir'),
+                gene_list_path=args.get('scfm_genes_list_path'),
+                gene_mapping_path=args.get('gene_mapping_path', '../gene2ens.csv'),
+                verbose=True
+            )
+        else:
+            raise ValueError("module_sets must be provided (via hyper_params or module_data_path/module_data_dir)")
+
+        # ============================================================
+        # [Ablation Experiment]shuffle_module_genes
+        # Disrupt the gene assignment within the module and verify the importance of the module structure (gene grouping)
+        # If performance decreases after scrambling, it means that the module structure contains useful biological information.
+        # ============================================================
+        if args.get('shuffle_module_genes', False):
+            shuffle_seed = args.get('shuffle_seed', 2021)
+            print(
+                f"[GEARS_Model_Module] Applying downstream-only extra module reshuffle "
+                f"(seed={shuffle_seed}). Pretraining-time module_sets from PKL will be reshuffled again."
+            )
+            all_gene_indices = []
+            module_sizes = []
+            for module in module_sets:
+                all_gene_indices.extend(module)
+                module_sizes.append(len(module))
+
+            import random
+            rng = random.Random(shuffle_seed)
+            rng.shuffle(all_gene_indices)
+
+            shuffled_module_sets = []
+            start_idx = 0
+            for size in module_sizes:
+                shuffled_module_sets.append(all_gene_indices[start_idx:start_idx + size])
+                start_idx += size
+            module_sets = shuffled_module_sets
+            print(f"[GEARS_Model_Module] Shuffled {len(all_gene_indices)} genes across {len(module_sets)} modules")
+
+        # ============================================================
+        # [Perturbation Embedding Layer]
+        # pert_w: Convert perturbation markers (0/1) to embedding vectors
+        # gene_emb: Basic embedding of genes (used when pre-trained models are not available)
+        # pert_emb: dedicated embedding of perturbed genes
+        # ============================================================
+        self.pert_w = nn.Linear(1, hidden_size)
+
+        # gene/global perturbation embedding dictionary lookup
+        self.gene_emb = nn.Embedding(self.num_genes, hidden_size, max_norm=True)
+        self.pert_emb = nn.Embedding(self.num_genes, hidden_size, max_norm=True)
+
+        # transformation layer
+        self.emb_trans = nn.ReLU()
+        self.pert_base_trans = nn.ReLU()
+        self.transform = nn.ReLU()
+        self.pert_fuse = MLP([hidden_size, hidden_size, hidden_size], last_layer_act='ReLU')
+
+        # ============================================================
+        # [GO graph neural network]
+        # G_sim: edge index of GO similarity graph
+        # G_sim_weight: edge weight
+        # Using SGConv (Simplified Graph Convolution) for message passing
+        # GO diagram helps the model understand gene functional relationships, and perturbation effects propagate along the diagram.
+        # ============================================================
+        self.G_sim = args['G_go'].to(args['device'])
+        self.G_sim_weight = args['G_go_weight'].to(args['device'])
+
+        self.sim_layers = torch.nn.ModuleList()
+        for i in range(1, self.num_layers + 1):
+            self.sim_layers.append(SGConv(hidden_size, hidden_size, 1))
+
+        # ============================================================
+        # [Decoder]
+        # recovery_w: shared MLP, projects embeddings into decoding space
+        # indv_w1/b1: Gene specific parameters, each gene has independent decoding weight
+        # cross_gene_state: cross-gene MLP, modeling inter-gene dependencies
+        # indv_w2/b2: final decoding layer
+        # ============================================================
+        self.recovery_w = MLP([hidden_size, hidden_size*2, hidden_size], last_layer_act='linear')
+
+        # gene specific decoder
+        self.indv_w1 = nn.Parameter(torch.rand(self.num_genes, hidden_size, 1))
+        self.indv_b1 = nn.Parameter(torch.rand(self.num_genes, 1))
+        self.act = nn.ReLU()
+        nn.init.xavier_normal_(self.indv_w1)
+        nn.init.xavier_normal_(self.indv_b1)
+
+        # Cross gene MLP
+        self.cross_gene_state = MLP([self.num_genes, hidden_size, hidden_size])
+        self.indv_w2 = nn.Parameter(torch.rand(1, self.num_genes, hidden_size+1))
+        self.indv_b2 = nn.Parameter(torch.rand(1, self.num_genes))
+        nn.init.xavier_normal_(self.indv_w2)
+        nn.init.xavier_normal_(self.indv_b2)
+
+        # BatchNorm layers
+        self.bn_emb = nn.BatchNorm1d(hidden_size)
+        self.bn_pert_base = nn.BatchNorm1d(hidden_size)
+        self.bn_pert_base_trans = nn.BatchNorm1d(hidden_size)
+
+        # uncertainty mode (optional)
+        if self.uncertainty:
+            self.uncertainty_w = MLP([hidden_size, hidden_size*2, hidden_size, 1], last_layer_act='linear')
+
+        # ============================================================
+        # [Core] Initialize pre-trained encoder (SCFM_Module)
+        # requires_grad: whether to fine-tune pre-training parameters
+        # - finetune_method='frozen': freeze pre-training parameters
+        # - Others: Full parameter fine-tuning
+        # ============================================================
+        self.use_module_level = args.get('use_module_level', False)
+        self.singlecell_model = SCFM_Module(
+            self.input_genes_list,
+            self.scfm_genes_list,
+            self.scfm_hyper_params,
+            self.scfm_ckpt,
+            module_sets=module_sets,
+            requires_grad=args.get('finetune_method') != 'frozen',
+            device=args['device'],
+            output_dim=hidden_size,
+            use_module_level=self.use_module_level
+        )
+        self.pretrained = True
+
+        # ============================================================
+        # [Experimental feature] module level features
+        # In the current implementation, module-level downstream tasks are not fully implemented.
+        # GNN and decoder need to be redesigned to handle module dimensions
+        # ============================================================
+        if self.use_module_level:
+            self.num_modules = len(module_sets)
+            import warnings
+            warnings.warn(
+                f"Using module-level features. This requires compatible downstream architecture. "
+                f"Number of modules: {self.num_modules}, Number of genes: {self.num_genes}. "
+                f"Currently module-level downstream task is not fully implemented.",
+                UserWarning
+            )
+
+    def forward(self, data):
+        """
+        Forward propagation: predicting expression values after gene perturbation
+
+        [enter]
+        data: PyG Data object
+            - x: (B*N, 2) characteristic matrix
+                [:, 0]: gene expression value (control group)
+                [:, 1]: perturbation flag (1 indicates that the gene is perturbed)
+            - y: real expression value (used during training)
+            - pert: disturbance type identifier
+
+        [Output]
+        out: (B, N) predicted gene expression value after perturbation
+
+        [Processing Process]
+        1. Obtain base gene embeddings (from pretrained encoder or random embeddings)
+        2. Generate perturbation embedding (enhanced by GO-GNN)
+        3. Merge base embedding and perturbation embedding
+        4. Predict the expression value through the decoder
+        """
+        if str(type(data)) == "<class 'torch.Tensor'>":
+            x = data.reshape((-1, 2))
+        else:
+            x, batch = data.x, data.batch
+        x_pert = x[:, 1].reshape((-1, self.num_genes))  # Disturbance mark
+        num_graphs = len(data.batch.unique())
+
+        # ============================================================
+        # [Step 1] Obtain basic gene embedding
+        # If using a pretrained model, get embeddings from gene expression encoding
+        # Otherwise use randomly initialized gene embeddings
+        # ============================================================
+        if self.pretrained:
+            exp = x[:, 0].reshape((-1, self.num_genes))  # B, N - gene expression value
+            emb = self.singlecell_model(exp)  # B, N, D - by pre-trained encoder
+            emb = emb.reshape((num_graphs * self.num_genes, -1))
+        else:
+            emb = self.gene_emb(torch.LongTensor(list(range(self.num_genes))).repeat(num_graphs, ).to(self.args['device']))
+        emb = self.bn_emb(emb)
+        base_emb = self.emb_trans(emb)
+
+        base_emb = base_emb.reshape(num_graphs, self.num_genes, -1)
+
+        pert = x_pert.reshape(-1, 1)  # B * N, 1
+
+        # ============================================================
+        # [Step 2] Generate perturbation embedding (GO-GNN enhancement)
+        # Perturbation embeddings are propagated through GO similarity graphs to capture gene functional relationships
+        # This enables the model to exploit a priori knowledge of gene function
+        # ============================================================
+        pert_embeddings = self.pert_emb(torch.LongTensor(list(range(self.num_genes))).to(self.args['device']))  # N, D
+
+        # Propagating perturbation embeddings via GO-GNN
+        for idx, layer in enumerate(self.sim_layers):
+            pert_embeddings = layer(pert_embeddings, self.G_sim, self.G_sim_weight)
+            if idx < self.num_layers - 1:
+                pert_embeddings = pert_embeddings.relu()
+
+        # After GNN: N, D
+        pert_embeddings = torch.stack(list((pert_embeddings,)) * num_graphs)   # B, N, D
+        pert_embeddings = pert_embeddings * x_pert.reshape((num_graphs, self.num_genes, 1))  # Only perturbed genes are retained
+        pert_embeddings = pert_embeddings.sum(1)  # B, D - aggregate all perturbed genes
+        pert_mask = (x_pert.sum(1) != 0).reshape((num_graphs, 1))  # B, 1
+        pert_embeddings = self.pert_fuse(pert_embeddings) * pert_mask  # B, D
+        pert_embeddings = pert_embeddings.reshape((num_graphs, 1, pert_embeddings.shape[1]))  # B, 1, D
+
+        # ============================================================
+        # [Step 3] Merge base embedding and perturbation embedding
+        # Perturbation information is added to the embedding of each gene
+        # ============================================================
+        base_emb = base_emb + pert_embeddings
+
+        base_emb = base_emb.reshape(num_graphs * self.num_genes, -1)
+
+        # Add perturbed position embedding
+        pert_emb = self.pert_w(pert)
+        combined = pert_emb + base_emb
+        combined = self.bn_pert_base_trans(combined)
+        base_emb = self.pert_base_trans(combined)
+        base_emb = self.bn_pert_base(base_emb)
+
+        # ============================================================
+        # [Step 4] Decoder: predict the expression value after perturbation
+        # Using gene-specific decoders + cross-gene MLP
+        # ============================================================
+        base_emb = self.transform(base_emb)
+        out = self.recovery_w(base_emb)  # B * N, D
+        out = out.reshape(num_graphs, self.num_genes, -1)  # B, N, D
+        out = out.unsqueeze(-1) * self.indv_w1  # B, N, D, 1 - Gene specific transformation
+        w = torch.sum(out, axis=2)  # B, N, 1
+        out = w + self.indv_b1  # B, N, 1
+
+        # Cross gene MLP: Modeling inter-gene dependencies
+        cross_gene_embed = self.cross_gene_state(out.reshape(num_graphs, self.num_genes, -1).squeeze(2))  # B, D
+        cross_gene_embed = cross_gene_embed.repeat(1, self.num_genes)  # B, N*D
+
+        cross_gene_embed = cross_gene_embed.reshape([num_graphs, self.num_genes, -1])  # B, N, D
+        cross_gene_out = torch.cat([out, cross_gene_embed], 2)  # B, N, D+1
+
+        cross_gene_out = cross_gene_out * self.indv_w2  # B, N, D+1
+        cross_gene_out = torch.sum(cross_gene_out, axis=2)  # B, N
+        out = cross_gene_out + self.indv_b2  # B, N
+
+        # ============================================================
+        # [Key] Residual connection: predicted value = change + original value
+        # This is the core design of GEARS. The model predicts expression changes rather than absolute values.
+        # ============================================================
+        out = out.reshape(num_graphs * self.num_genes, -1) + x[:, 0].reshape(-1, 1)  # B * N, 1
+
+        # ============================================================
+        # [Prediction record] for subsequent analysis
+        # ============================================================
+        if self.args.get('start_record', False):
+            y_flat = data.y.reshape(-1, 1)  # B * N, 1
+            pred = torch.cat((pert, out, y_flat), dim=1)  # B * N, 3
+            pred = pred.cpu().to_sparse()
+
+            values = pred.values().numpy()
+            indices = pred.indices().numpy()
+            sparse_matrix = scipy.sparse.coo_matrix((values, indices), shape=pred.shape)
+
+            scipy.sparse.save_npz(f'{self.pred_dir}/b{self.pred_batch_idx}_{num_graphs}.npz', sparse_matrix)
+            self.pred_batch_idx += 1
+
+        # ============================================================
+        # [Key] reshape output is (B, N) format
+        # In this way, the loss function can correctly index the predicted value and the true value.
+        # ============================================================
+        out = torch.split(torch.flatten(out), self.num_genes)
+        out = torch.stack(out)  # (B, N)
+
+        # uncertainty mode (optional)
+        if self.uncertainty:
+            logvar = self.uncertainty_w(base_emb)
+            logvar = torch.split(torch.flatten(logvar), self.num_genes)
+            return out, torch.stack(logvar)
+
+        return out  # (B, N)
+

BioMod/gene_perturbation/pertdata_local.py

@@ -0,0 +1,524 @@
+"""
+================================================================================
+Perturbation data processing module (pertdata_local.py)
+================================================================================
+
+[Module function]
+Load and preprocess the genetic perturbation dataset, converting the raw data into a model-usable PyG graph format.
+
+[Supported data sets]
+- norman: Norman et al. 2019, Large-scale CRISPR perturbation data
+- adamson: Adamson et al. 2016, single gene perturbation data
+- dixit: Dixit et al. 2016, single gene perturbation data
+- replogle_k562_essential/replogle_rpe1_essential: Replogle et al. 2022
+
+[Data format]
+Each cell is represented as a graph data object:
+- x: (num_genes, 2) feature matrix, [:, 0] is the expression value, [:, 1] is the perturbation mark
+- y: target expression value
+- de_idx: differentially expressed gene index
+- pert: disturbance type identifier
+
+[Key steps]
+1. Load the original h5ad file
+2. Calculate differentially expressed genes
+3. Divide training/validation/test sets
+4. Create PyG graph data objects
+================================================================================
+"""
+
+from torch_geometric.data import Data
+import torch
+import numpy as np
+import pickle
+from torch_geometric.data import DataLoader
+import os
+import scanpy as sc
+import networkx as nx
+from tqdm import tqdm
+import scipy.sparse as sp
+
+import warnings
+warnings.filterwarnings("ignore")
+sc.settings.verbosity = 0
+
+# ============================================================
+# Use local modules and avoid external gears_reduce imports
+# ============================================================
+from utils import print_sys, zip_data_download_wrapper, filter_pert_cond_in_go
+from data_utils_local import get_DE_genes, get_dropout_non_zero_genes, DataSplitter
+
+class PertData:
+
+    def __init__(self, data_path):
+        """
+        Initialize the perturbation data processing class
+
+        Args:
+            data_path: data storage directory path
+        """
+        self.data_path = data_path
+        self.gene_names = None
+        if not os.path.exists(self.data_path):
+            os.mkdir(self.data_path)
+            
+    def load(self, data_name = None,
+             data_path = None, exclude_pert = set()):
+        """
+        Load perturbation dataset
+
+        [Data source]
+        - Download public data sets such as norman/adamson/dixit from Harvard Dataverse
+        - You can also provide the local h5ad file path
+
+        [Processing Process]
+        1. Download or load data
+        2. Filter disturbances that are not in the GO graph
+        3. Create PyG graph data objects
+        """
+        if data_name in ['norman', 'adamson', 'dixit',
+                         'replogle_k562_essential',
+                         'replogle_rpe1_essential']:
+            # Download public datasets from Harvard Dataverse
+            if data_name == 'norman':
+                url = 'https://dataverse.harvard.edu/api/access/datafile/6154020'
+            elif data_name == 'adamson':
+                url = 'https://dataverse.harvard.edu/api/access/datafile/6154417'
+            elif data_name == 'dixit':
+                url = 'https://dataverse.harvard.edu/api/access/datafile/6154416'
+            elif data_name == 'replogle_k562_essential':
+                url = 'https://dataverse.harvard.edu/api/access/datafile/7458695'
+            elif data_name == 'replogle_rpe1_essential':
+                url = 'https://dataverse.harvard.edu/api/access/datafile/7458694'
+            data_path = os.path.join(self.data_path, data_name)
+            zip_data_download_wrapper(url, data_path, self.data_path)
+            self.dataset_name = data_path.split('/')[-1]
+            self.dataset_path = data_path
+            adata_path = os.path.join(data_path, 'perturb_processed.h5ad')
+            self.adata = sc.read_h5ad(adata_path)
+
+        elif os.path.exists(data_path):
+            # Load local data
+            adata_path = os.path.join(data_path, 'perturb_processed.h5ad')
+            self.adata = sc.read_h5ad(adata_path)
+            self.dataset_name = data_path.split('/')[-1]
+            self.dataset_path = data_path
+        else:
+            raise ValueError("data is either Norman/Adamson/Dixit or a path to an h5ad file")
+
+        # Filter perturbed genes that are not in the GO graph
+        print_sys('These perturbations are not in the GO graph and their '
+                  'perturbation can thus not be predicted')
+        print_sys(exclude_pert)
+        if len(exclude_pert) > 0:
+            filter_go = self.adata.obs[self.adata.obs.condition.apply(
+                                lambda x: filter_pert_cond_in_go(x, exclude_pert))]
+            self.adata = self.adata[filter_go.index.values, :]
+
+        # ============================================================
+        # [Key] Create or load PyG graph data objects
+        # This is the core step of data preprocessing, converting each cell into graph format
+        # ============================================================
+        pyg_path = os.path.join(data_path, 'data_pyg')
+        if not os.path.exists(pyg_path):
+            os.mkdir(pyg_path)
+        dataset_fname = os.path.join(pyg_path, 'cell_graphs.pkl')
+
+        if os.path.isfile(dataset_fname):
+            # Already processed data can be loaded directly.
+            print_sys("Local copy of pyg dataset is detected. Loading...")
+            self.dataset_processed = pickle.load(open(dataset_fname, "rb"))
+            print_sys("Done!")
+        else:
+            # Data needs to be reprocessed
+            self.ctrl_adata = self.adata[self.adata.obs['condition'] == 'ctrl']
+            self.gene_names = list(self.adata.var.gene_name)
+            conditions = list(self.adata.obs.condition.unique())
+            pert_genes = []
+            for condition in conditions:
+                if condition == 'ctrl':
+                    continue
+                # Analyze perturbation combinations (format: gene1+gene2)
+                cond1 = condition.split('+')[0]
+                cond2 = condition.split('+')[1]
+                if cond1 != 'ctrl':
+                    if cond1 not in pert_genes:
+                        pert_genes.append(cond1)
+                    if cond1 not in self.gene_names:
+                        self.gene_names.append(cond1)
+                if cond2 != 'ctrl':
+                    if cond2 not in pert_genes:
+                        pert_genes.append(cond2)
+                    if cond2 not in self.gene_names:
+                        self.gene_names.append(cond2)
+            self.gene_names = np.array(self.gene_names)
+            self.real_gene_names = np.array(list(self.adata.var.gene_name))
+            print_sys(f"len(self.gene_names): {len(self.gene_names)}")
+            print_sys(f"len(self.real_gene_names): {len(self.real_gene_names)}")
+            print_sys("Creating pyg object for each cell in the data...")
+            self.dataset_processed = self.create_dataset_file()
+            print_sys("Saving new dataset pyg object at " + dataset_fname)
+            pickle.dump(self.dataset_processed, open(dataset_fname, "wb"))
+            print_sys("Done!")
+            
+    def new_data_process(self, dataset_name,
+                         adata = None):
+        
+        if 'condition' not in adata.obs.columns.values:
+            raise ValueError("Please specify condition")
+        if 'gene_name' not in adata.var.columns.values:
+            raise ValueError("Please specify gene name")
+        if 'cell_type' not in adata.obs.columns.values:
+            raise ValueError("Please specify cell type")
+        
+        dataset_name = dataset_name.lower()
+        self.dataset_name = dataset_name
+        save_data_folder = os.path.join(self.data_path, dataset_name)
+        
+        if not os.path.exists(save_data_folder):
+            os.mkdir(save_data_folder)
+        self.dataset_path = save_data_folder
+        self.adata = get_DE_genes(adata)
+        self.adata = get_dropout_non_zero_genes(self.adata)
+        self.adata.write_h5ad(os.path.join(save_data_folder, 'perturb_processed.h5ad'))
+        
+        self.ctrl_adata = self.adata[self.adata.obs['condition'] == 'ctrl']
+    
+        self.gene_names = list(self.adata.var.gene_name)
+        conditions = list(self.adata.obs.condition.unique())
+        pert_genes = []
+        for condition in conditions:
+            if condition == 'ctrl':
+                continue
+            cond1 = condition.split('+')[0]
+            cond2 = condition.split('+')[1]
+            if cond1 != 'ctrl': 
+                if cond1 not in pert_genes:
+                    pert_genes.append(cond1)
+                if cond1 not in self.gene_names:
+                    self.gene_names.append(cond1)
+            if cond2 != 'ctrl':
+                if cond2 not in pert_genes:
+                    pert_genes.append(cond2)
+                if cond2 not in self.gene_names:
+                    self.gene_names.append(cond2)
+        self.gene_names = np.array(self.gene_names)
+        self.real_gene_names = np.array(list(self.adata.var.gene_name))
+
+
+        pyg_path = os.path.join(save_data_folder, 'data_pyg')
+        if not os.path.exists(pyg_path):
+            os.mkdir(pyg_path)
+        dataset_fname = os.path.join(pyg_path, 'cell_graphs.pkl')
+        print_sys("Creating pyg object for each cell in the data...")
+        self.dataset_processed = self.create_dataset_file()
+        print_sys("Saving new dataset pyg object at " + dataset_fname) 
+        pickle.dump(self.dataset_processed, open(dataset_fname, "wb"))    
+        print_sys("Done!")
+        
+    def prepare_split(self, split = 'simulation',
+                      seed = 1,
+                      train_gene_set_size = 0.75,
+                      combo_seen2_train_frac = 0.75,
+                      combo_single_split_test_set_fraction = 0.1,
+                      test_perts = None,
+                      only_test_set_perts = False,
+                      test_pert_genes = None):
+        """
+        Prepare data partitioning
+
+        [Description of partition strategy]
+        - 'simulation': simulate the division of real application scenarios
+            The test set contains: unseen_single (unseen single gene perturbation),
+                       combo_seen0/1/2 (combination perturbations with different familiarity levels)
+            This is the most commonly used division method to evaluate the generalization ability of the model.
+        - 'combo_seen0/1/2': divided according to the number of genes seen in the combination
+        - 'single': only evaluate single gene perturbations
+        - 'no_test': Do not divide the test set
+
+        [Key parameters]
+        - train_gene_set_size: training set gene proportion (0.75 means 75% of genes are used for training)
+        - seed: random seed to ensure the experiment is reproducible
+        """
+        available_splits = ['simulation', 'simulation_single', 'combo_seen0', 'combo_seen1',
+                            'combo_seen2', 'single', 'no_test', 'no_split']
+        if split not in available_splits:
+            raise ValueError('currently, we only support ' + ','.join(available_splits))
+        self.split = split
+        self.seed = seed
+        self.subgroup = None
+        self.train_gene_set_size = train_gene_set_size
+
+        # Divide file path
+        split_folder = os.path.join(self.dataset_path, 'splits')
+        if not os.path.exists(split_folder):
+            os.mkdir(split_folder)
+        split_file = self.dataset_name + '_' + split + '_' + str(seed) + '_' + str(train_gene_set_size) + '.pkl'
+        split_path = os.path.join(split_folder, split_file)
+
+        if test_perts:
+            split_path = split_path[:-4] + '_' + test_perts + '.pkl'
+
+        if os.path.exists(split_path):
+            # Already divided, load directly
+            print_sys("Local copy of split is detected. Loading...")
+            set2conditions = pickle.load(open(split_path, "rb"))
+            if split == 'simulation':
+                subgroup_path = split_path[:-4] + '_subgroup.pkl'
+                subgroup = pickle.load(open(subgroup_path, "rb"))
+                self.subgroup = subgroup
+        else:
+            # Create new partition
+            print_sys("Creating new splits....")
+            if test_perts:
+                test_perts = test_perts.split('_')
+
+            if split in ['simulation', 'simulation_single']:
+                DS = DataSplitter(self.adata, split_type=split)
+
+                adata, subgroup = DS.split_data(train_gene_set_size = train_gene_set_size,
+                                                combo_seen2_train_frac = combo_seen2_train_frac,
+                                                seed=seed,
+                                                test_perts = test_perts,
+                                                only_test_set_perts = only_test_set_perts
+                                               )
+                subgroup_path = split_path[:-4] + '_subgroup.pkl'
+                pickle.dump(subgroup, open(subgroup_path, "wb"))
+                self.subgroup = subgroup
+
+            elif split[:5] == 'combo':
+                split_type = 'combo'
+                seen = int(split[-1])
+
+                if test_pert_genes:
+                    test_pert_genes = test_pert_genes.split('_')
+
+                DS = DataSplitter(self.adata, split_type=split_type, seen=int(seen))
+                adata = DS.split_data(test_size=combo_single_split_test_set_fraction,
+                                      test_perts=test_perts,
+                                      test_pert_genes=test_pert_genes,
+                                      seed=seed)
+
+            elif split == 'single':
+                DS = DataSplitter(self.adata, split_type=split)
+                adata = DS.split_data(test_size=combo_single_split_test_set_fraction, seed=seed)
+
+            elif split == 'no_test':
+                DS = DataSplitter(self.adata, split_type=split)
+                adata = DS.split_data(test_size=combo_single_split_test_set_fraction, seed=seed)
+
+            elif split == 'no_split':
+                adata = self.adata
+                adata.obs['split'] = 'test'
+
+            # Save partition results
+            set2conditions = dict(adata.obs.groupby('split').agg({'condition': lambda x: x}).condition)
+            set2conditions = {i: j.unique().tolist() for i,j in set2conditions.items()}
+            pickle.dump(set2conditions, open(split_path, "wb"))
+            print_sys("Saving new splits at " + split_path)
+
+        self.set2conditions = set2conditions
+
+        if split == 'simulation':
+            print_sys('Simulation split test composition:')
+            for i,j in subgroup['test_subgroup'].items():
+                print_sys(i + ':' + str(len(j)))
+        print_sys("Done!")
+        
+    def get_dataloader(self, batch_size, test_batch_size = None):
+        """
+        Create data loader
+
+        [Data loader]
+        Encapsulate PyG graph data objects as DataLoader, supporting batch processing and shuffling.
+
+        [Key settings]
+        - Training set: shuffle=True, drop_last=True (make sure each batch size is consistent)
+        - Validation/test set: shuffle=False (maintain order for easy analysis)
+        """
+        if test_batch_size is None:
+            test_batch_size = batch_size
+        if self.gene_names is None:
+            self.gene_names = list(self.adata.var.gene_name)
+            conditions = list(self.adata.obs.condition.unique())
+            pert_genes = []
+            for condition in conditions:
+                if condition == 'ctrl':
+                    continue
+                cond1 = condition.split('+')[0]
+                cond2 = condition.split('+')[1]
+                if cond1 != 'ctrl':
+                    if cond1 not in pert_genes:
+                        pert_genes.append(cond1)
+                    if cond1 not in self.gene_names:
+                        self.gene_names.append(cond1)
+                if cond2 != 'ctrl':
+                    if cond2 not in pert_genes:
+                        pert_genes.append(cond2)
+                    if cond2 not in self.gene_names:
+                        self.gene_names.append(cond2)
+            self.gene_names = np.array(self.gene_names)
+            self.real_gene_names = np.array(list(self.adata.var.gene_name))
+
+        # Gene name to index mapping
+        self.node_map = {x: it for it, x in enumerate(self.gene_names)}
+        if (self.gene_names[: len(self.real_gene_names)] != self.real_gene_names).any():
+            raise ValueError("self.real_gene_names must be in the start of self.gene_names")
+
+        # Create cytogram data
+        cell_graphs = {}
+        if self.split == 'no_split':
+            i = 'test'
+            cell_graphs[i] = []
+            for p in self.set2conditions[i]:
+                if p != 'ctrl':
+                    cell_graphs[i].extend(self.dataset_processed[p])
+
+            print_sys("Creating dataloaders....")
+            test_loader = DataLoader(cell_graphs['test'],
+                                batch_size=batch_size, shuffle=False)
+
+            print_sys("Dataloaders created...")
+            return {'test_loader': test_loader}
+        else:
+            if self.split =='no_test':
+                splits = ['train','val']
+            else:
+                splits = ['train','val','test']
+            for i in splits:
+                cell_graphs[i] = []
+                for p in self.set2conditions[i]:
+                    cell_graphs[i].extend(self.dataset_processed[p])
+
+            print_sys("Creating dataloaders....")
+
+            # Set up data loader
+            train_loader = DataLoader(cell_graphs['train'],
+                                batch_size=batch_size, shuffle=True, drop_last = True)
+            val_loader = DataLoader(cell_graphs['val'],
+                                batch_size=batch_size, shuffle=True)
+
+            if self.split !='no_test':
+                test_loader = DataLoader(cell_graphs['test'],
+                                batch_size=batch_size, shuffle=False)
+                self.dataloader =  {'train_loader': train_loader,
+                                    'val_loader': val_loader,
+                                    'test_loader': test_loader}
+
+            else:
+                self.dataloader =  {'train_loader': train_loader,
+                                    'val_loader': val_loader}
+            print_sys("Done!")
+        del self.dataset_processed  # Clean memory
+    
+        
+    def create_dataset_file(self):
+        dl = {}
+        for p in tqdm(self.adata.obs['condition'].unique()):
+            cell_graph_dataset = self.create_cell_graph_dataset(self.adata, p, num_samples=1)
+            dl[p] = cell_graph_dataset
+        return dl
+    
+    def get_pert_idx(self, pert_category, adata_):
+        pert_idx = [np.where(p == self.gene_names)[0][0]
+                    for p in pert_category.split('+')
+                    if p != 'ctrl']
+        return pert_idx
+
+    # Set up feature matrix and output
+    def create_cell_graph(self, X, y, de_idx, pert, pert_idx=None):
+        """
+        Create a graph data object of a single cell
+
+        [enter]
+        - X: Control group gene expression value (expression before perturbation)
+        - y: gene expression value after perturbation (target value)
+        - de_idx: index of differentially expressed genes
+        - pert: disturbance type identifier
+        - pert_idx: index list of perturbed genes
+
+        [Output]
+        PyG Data object:
+        - x: (num_genes, 2) feature matrix
+            [:, 0]: Control group expression value
+            [:, 1]: perturbation flag (1=the gene is perturbed)
+        - y: target expression value
+        - de_idx: differentially expressed gene index
+        - pert: perturbation type
+        """
+        # Create perturbation feature vectors (mark which genes are perturbed)
+        pert_feats = np.zeros(len(X[0]))
+        if pert_idx is not None:
+            for p in pert_idx:
+                pert_feats[int(np.abs(p))] = 1
+        pert_feats = np.expand_dims(pert_feats, 0)
+        # Spliced expression values and perturbation markers
+        feature_mat = torch.Tensor(np.concatenate([X, pert_feats])).T
+
+        return Data(x=feature_mat, edge_index=None, edge_attr=None,
+                    y=torch.Tensor(y), de_idx=de_idx, pert=pert)
+
+    def create_cell_graph_dataset(self, split_adata, pert_category,
+                                  num_samples=1):
+        """
+        Create a graph dataset of perturbation categories
+
+        [Processing logic]
+        1. For non-control group disturbances:
+           - Get perturbed gene index
+           - Expression before perturbation = randomly sample control cells
+           - Target = expression of cells after perturbation
+        2. For the control group:
+           - Expression before perturbation = target = expression in control cells
+
+        [num_samples parameter]
+        Each perturbed cell can be paired with multiple control group cells,
+        Increase the diversity of training data
+        """
+        num_de_genes = 20  # Only focus on the top 20 differentially expressed genes for evaluation
+        adata_ = split_adata[split_adata.obs['condition'] == pert_category]
+        de_genes = adata_.uns['rank_genes_groups_cov_all']
+        Xs = []
+        ys = []
+
+        # Handling non-control group disturbances
+        if pert_category != 'ctrl':
+            # Get the index of the perturbed gene
+            pert_idx = self.get_pert_idx(pert_category, adata_)
+
+            # Get differentially expressed gene index
+            pert_de_category = adata_.obs['condition_name'][0]
+            de_idx = np.where(adata_.var_names.isin(
+                np.array(de_genes[pert_de_category][:num_de_genes])))[0]
+
+            # Expand feature matrix to include perturbed genes
+            padded_X = sp.lil_matrix((adata_.X.shape[0], len(self.gene_names)), dtype=adata_.X.dtype)
+            padded_X[:, :len(self.real_gene_names)] = adata_.X
+
+            for cell_z in padded_X:
+                # Use control group samples as expression before perturbation
+                ctrl_samples = self.ctrl_adata[np.random.randint(0,
+                                        len(self.ctrl_adata), num_samples), :]
+                padded_ctrl_X = sp.lil_matrix((ctrl_samples.X.shape[0], len(self.gene_names)), dtype=ctrl_samples.X.dtype)
+                padded_ctrl_X[:, :len(self.real_gene_names)] = ctrl_samples.X
+                for c in padded_ctrl_X:
+                    Xs.append(c)  # control group expression
+                    ys.append(cell_z)  # expression after perturbation
+
+        # Process control group
+        else:
+            pert_idx = None
+            de_idx = [-1] * num_de_genes
+            padded_X = sp.lil_matrix((adata_.X.shape[0], len(self.gene_names)), dtype=adata_.X.dtype)
+            padded_X[:, :len(self.real_gene_names)] = adata_.X
+            for cell_z in padded_X:
+                Xs.append(cell_z)
+                ys.append(cell_z)
+
+        # Create a cell diagram data object
+        cell_graphs = []
+        for X, y in zip(Xs, ys):
+            cell_graphs.append(self.create_cell_graph(X.toarray(),
+                                y.toarray(), de_idx, pert_category, pert_idx))
+
+        return cell_graphs

BioMod/gene_perturbation/scfm_utils_local.py

@@ -0,0 +1,777 @@
+"""
+================================================================================
+Single cell basic model tool module (scfm_utils_local.py)
+================================================================================
+
+[Module function]
+Provides basic tool functions for distributed training and model management, including:
+
+1. Distributed training tools:
+   - MetricLogger: training process indicator record
+   - SequentialDistributedSampler: Distributed sequential sampler
+   - get_reduced/distributed_concat: cross-process data aggregation
+
+2. Learning rate scheduler:
+   - CosineAnnealingWarmupRestarts: Cosine Annealing with warm-up and restart
+
+3. Model saving/loading:
+   - save_ckpt/save_scaler_ckpt: save checkpoint
+   - load_model: Load pre-trained model
+
+[Notes on distributed training]
+- RANK, WORLD_SIZE, LOCAL_RANK environment variables need to be set correctly
+- The main process (rank=0) is responsible for printing logs and saving models
+- Data needs to be aggregated across processes to get global indicators
+================================================================================
+"""
+
+# -*- coding: utf-8 -*-
+
+from __future__ import print_function
+import json
+import os
+import struct
+import sys
+import platform
+import re
+import time
+import traceback
+import requests
+import socket
+import random
+import math
+import numpy as np
+import torch
+import logging
+import datetime
+from torch.optim.lr_scheduler import _LRScheduler
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.modules.loss import _WeightedLoss
+
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch import inf
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print('[{}] '.format(now), end='')  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+        args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+        args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}, gpu {}'.format(
+        args.rank, args.dist_url, args.gpu), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % args.resume)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x
+
+
+
+def seed_all(seed_value, cuda_deterministic=False):
+    """
+    Set all random seeds
+    """
+    random.seed(seed_value)
+    os.environ['PYTHONHASHSEED'] = str(seed_value)
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed_value)
+        torch.cuda.manual_seed_all(seed_value)
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    if cuda_deterministic:  # slower, more reproducible
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:  # faster, less reproducible
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+
+
+def set_log(logfileName, rank=-1):
+    """
+    The master node saves all logs, and other nodes only save warnings and errors.
+    """
+    log_file_folder = os.path.dirname(logfileName)
+    time_now = datetime.datetime.now()
+    logfileName = f'{logfileName}_{time_now.year}_{time_now.month}_{time_now.day}_{time_now.hour}_{time_now.minute}.log'
+    if not os.path.exists(log_file_folder):
+        os.makedirs(log_file_folder)
+    else:
+        pass
+
+    logging.basicConfig(level=logging.INFO if rank in [-1, 0] else logging.WARN,
+        format='[%(asctime)s %(levelname)s %(filename)s line %(lineno)d %(process)d] %(message)s',
+        datefmt='[%X]',
+        handlers=[logging.FileHandler(logfileName), logging.StreamHandler()]
+    )
+    logger = logging.getLogger()
+    return logger
+
+
+def save_ckpt(epoch, model, optimizer, scheduler, losses, model_name, ckpt_folder):
+    """
+    Save model checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    torch.save(
+        {
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'losses': losses,
+        },
+        f'{ckpt_folder}{model_name}_{epoch}.pth'
+    )
+
+
+def save_scaler_ckpt(epoch, model, optimizer, scheduler, scaler, losses, model_name, ckpt_folder, torch_compile = False):
+    """
+    Save model checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    if torch_compile:
+        model = model._orig_mod
+    torch.save(
+        {
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'scaler_state_dict': scaler.state_dict(),
+            'losses': losses,
+        },
+        f'{ckpt_folder}{model_name}_{epoch}.pth'
+    )
+
+def save_simple_ckpt(model, model_name, ckpt_folder):
+    """
+    Save model checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    torch.save(
+        {
+            'model_state_dict': model.module.state_dict()
+        },
+        f'{ckpt_folder}{model_name}.pth'
+    )
+
+def save_best_ckpt(epoch, model, optimizer, scheduler, losses, model_name, ckpt_folder):
+    """
+    Save model checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    torch.save(
+        {
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'losses': losses,
+        },
+        f'{ckpt_folder}{model_name}_best.pth'
+    )
+
+def get_reduced(tensor, current_device, dest_device, world_size):
+    """
+    Concentrate variables or tensors on different GPUs on the main GPU and get the mean
+    """
+    tensor = tensor.clone().detach() if torch.is_tensor(tensor) else torch.tensor(tensor)
+    tensor = tensor.to(current_device)
+    torch.distributed.reduce(tensor, dst=dest_device)
+    tensor_mean = tensor.item() / world_size
+    return tensor_mean
+
+def get_ndtensor_reduced(tensor, current_device, dest_device, world_size):
+    """
+    Concentrate variables or tensors on different GPUs on the main GPU and get the mean, which needs to be a 2-dimensional tensor.
+    """
+    tensor = tensor.clone().detach() if torch.is_tensor(tensor) else torch.tensor(tensor)
+    tensor = tensor.to(current_device)
+    torch.distributed.reduce(tensor, dst=dest_device)
+    tensor_mean = torch.zeros(tensor.shape)
+    if len(tensor.shape) == 2:
+        for i in range(tensor.shape[0]):
+            for j in range(tensor.shape[1]):
+                tensor_mean[i,j] = tensor[i,j].item() / world_size
+    elif len(tensor.shape) == 1:
+        for i in range(tensor.shape[0]):
+            tensor_mean[i] = tensor[i].item() / world_size
+    return tensor_mean
+
+def numel(m: torch.nn.Module, only_trainable: bool = False):
+    """
+    returns the total number of parameters used by `m` (only counting
+    shared parameters once); if `only_trainable` is True, then only
+    includes parameters with `requires_grad = True`
+    """
+    parameters = m.parameters()
+    if only_trainable:
+        parameters = list(p for p in parameters if p.requires_grad)
+    unique = dict((p.data_ptr(), p) for p in parameters).values()
+    return sum(p.numel() for p in unique)
+
+
+def label_smooth(y, K, epsilon=0.1):
+    """
+    Label smoothing for multiclass labels
+    One hot encode labels `y` over `K` classes. `y` should be of the form [1, 6, 3, etc.]
+    """
+    m = len(y)
+    out = np.ones((m, K)) * epsilon / K
+    for index in range(m):
+        out[index][y[index] - 1] += 1 - epsilon
+    return torch.tensor(out)
+
+
+class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):
+    """
+    sequential distributed sampler
+
+    [Function]
+    Sequentially subsample data in a distributed environment to facilitate post-training
+    Collect the results of all processes.
+
+    [Why is this needed?]
+    - Ordinary DistributedSampler will disrupt the order of data
+    - The order needs to be maintained during verification/testing to facilitate subsequent aggregation of results.
+    - Ensure equal number of samples per process by adding padding samples
+    """
+
+    def __init__(self, dataset, batch_size, world_size, rank=None, num_replicas=None):
+        if num_replicas is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = world_size
+        if rank is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = torch.distributed.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.batch_size = batch_size
+        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        indices = list(range(len(self.dataset)))
+        # Add padding samples to make the data size divisible
+        indices += [indices[-1]] * (self.total_size - len(indices))
+        # Subsampling: Each process only takes the part it is responsible for
+        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+
+def distributed_concat(tensor, num_total_examples, world_size):
+    """
+    Merging inference results from different processes
+
+    [Function]
+    Use all_gather to collect tensors from all processes, and then splice and truncate the filled samples.
+
+    [parameter]
+    - tensor: the result tensor of the current process
+    - num_total_examples: real total number of samples
+    - world_size: number of processes
+    """
+    output_tensors = [tensor.clone() for _ in range(world_size)]
+    torch.distributed.all_gather(output_tensors, tensor)
+    concat = torch.cat(output_tensors, dim=0)
+    # Truncate filler samples
+    return concat[:num_total_examples]
+
+
+class CosineAnnealingWarmupRestarts(_LRScheduler):
+    """
+        optimizer (Optimizer): Wrapped optimizer.
+        first_cycle_steps (int): First cycle step size.
+        cycle_mult(float): Cycle steps magnification. Default: -1.
+        max_lr(float): First cycle's max learning rate. Default: 0.1.
+        min_lr(float): Min learning rate. Default: 0.001.
+        warmup_steps(int): Linear warmup step size. Default: 0.
+        gamma(float): Decrease rate of max learning rate by cycle. Default: 1.
+        last_epoch (int): The index of last epoch. Default: -1.
+    """
+    
+    def __init__(self,
+                 optimizer : torch.optim.Optimizer,
+                 first_cycle_steps : int,
+                 cycle_mult : float = 1.,
+                 max_lr : float = 0.1,
+                 min_lr : float = 0.001,
+                 warmup_steps : int = 0,
+                 gamma : float = 1.,
+                 last_epoch : int = -1
+        ):
+        assert warmup_steps < first_cycle_steps
+        
+        self.first_cycle_steps = first_cycle_steps # first cycle step size
+        self.cycle_mult = cycle_mult # cycle steps magnification
+        self.base_max_lr = max_lr # first max learning rate
+        self.max_lr = max_lr # max learning rate in the current cycle
+        self.min_lr = min_lr # min learning rate
+        self.warmup_steps = warmup_steps # warmup step size
+        self.gamma = gamma # decrease rate of max learning rate by cycle
+        
+        self.cur_cycle_steps = first_cycle_steps # first cycle step size
+        self.cycle = 0 # cycle count
+        self.step_in_cycle = last_epoch # step size of the current cycle
+        
+        super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch)
+        
+        # set learning rate min_lr
+        self.init_lr()
+
+    def init_lr(self):
+        self.base_lrs = []
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = self.min_lr
+            self.base_lrs.append(self.min_lr)
+    
+    def get_lr(self):
+        if self.step_in_cycle == -1:
+            return self.base_lrs
+        elif self.step_in_cycle < self.warmup_steps:
+            return [(self.max_lr - base_lr)*self.step_in_cycle / self.warmup_steps + base_lr for base_lr in self.base_lrs]
+        else:
+            return [base_lr + (self.max_lr - base_lr) \
+                    * (1 + math.cos(math.pi * (self.step_in_cycle-self.warmup_steps) \
+                                    / (self.cur_cycle_steps - self.warmup_steps))) / 2
+                    for base_lr in self.base_lrs]
+
+    def step(self, epoch=None):
+        if epoch is None:
+            epoch = self.last_epoch + 1
+            self.step_in_cycle = self.step_in_cycle + 1
+            if self.step_in_cycle >= self.cur_cycle_steps:
+                self.cycle += 1
+                self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps
+                self.cur_cycle_steps = int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult) + self.warmup_steps
+        else:
+            if epoch >= self.first_cycle_steps:
+                if self.cycle_mult == 1.:
+                    self.step_in_cycle = epoch % self.first_cycle_steps
+                    self.cycle = epoch // self.first_cycle_steps
+                else:
+                    n = int(math.log((epoch / self.first_cycle_steps * (self.cycle_mult - 1) + 1), self.cycle_mult))
+                    self.cycle = n
+                    self.step_in_cycle = epoch - int(self.first_cycle_steps * (self.cycle_mult ** n - 1) / (self.cycle_mult - 1))
+                    self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (n)
+            else:
+                self.cur_cycle_steps = self.first_cycle_steps
+                self.step_in_cycle = epoch
+                
+        self.max_lr = self.base_max_lr * (self.gamma**self.cycle)
+        self.last_epoch = math.floor(epoch)
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+
+
+class DistanceLoss(_WeightedLoss):
+    """
+    CrossEntropyLoss with Distance Weighted
+    """
+    def __init__(self, weight=None, reduction='mean', ignore_index = None):
+        super().__init__(weight=weight, reduction=reduction)
+        self.weight = weight
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+    def forward(self, inputs, targets):
+        if len(inputs.shape) > 2:
+            inputs = inputs.reshape(-1, inputs.size(-1))
+        if len(targets.shape) > 1:
+            targets = targets.reshape(-1)
+        if self.ignore_index is not None:
+            keep_index = (targets != self.ignore_index).nonzero(as_tuple=True)[0]
+            targets = torch.index_select(targets, 0, keep_index) #targets[targets != self.ignore_index]
+            inputs = torch.index_select(inputs, 0, keep_index)
+        lsm = F.log_softmax(inputs, -1)
+        targets = torch.empty(size=(targets.size(0), inputs.size(-1)), device=targets.device).fill_(0).scatter_(1, targets.data.unsqueeze(1), 1)
+        if self.weight is not None:
+            lsm = lsm * self.weight.unsqueeze(0)
+        loss = -(targets * lsm).sum(-1)
+        inputs = nn.Softmax(dim=-1)(inputs)[..., 1:-1].argmax(dim=-1) + 1
+        # print('inputs', inputs.device, inputs.shape)
+        targets = nn.Softmax(dim=-1)(targets)[..., 1:-1].argmax(dim=-1) + 1
+        # print('targets', targets.device, targets.shape)
+        distance = abs(inputs - targets) + 1e-2
+        # print('loss.shape', loss.shape)
+        # print('distance.shape', distance.shape)
+        loss = loss * distance
+        if self.reduction == 'sum':
+            loss = loss.sum()
+        elif self.reduction == 'mean':
+            loss = loss.mean()
+        return loss
+
+
+class LabelSmoothCrossEntropyLoss(_WeightedLoss):
+    """
+    CrossEntropyLoss with Label Somoothing
+    """
+    def __init__(self, weight=None, reduction='mean', smoothing=0.0):
+        super().__init__(weight=weight, reduction=reduction)
+        self.smoothing = smoothing
+        self.weight = weight
+        self.reduction = reduction
+
+    @staticmethod
+    def _smooth_one_hot(targets: torch.Tensor, n_classes: int, smoothing=0.0):
+        assert 0 <= smoothing < 1
+        with torch.no_grad():
+            targets = torch.empty(size=(targets.size(0), n_classes),
+                                  device=targets.device) \
+                .fill_(smoothing / (n_classes - 1)) \
+                .scatter_(1, targets.data.unsqueeze(1), 1. - smoothing)
+        return targets
+
+    def forward(self, inputs, targets):
+        targets = LabelSmoothCrossEntropyLoss._smooth_one_hot(targets, inputs.size(-1),
+                                                              self.smoothing)
+        lsm = F.log_softmax(inputs, -1)
+
+        if self.weight is not None:
+            lsm = lsm * self.weight.unsqueeze(0)
+
+        loss = -(targets * lsm).sum(-1)
+
+        if self.reduction == 'sum':
+            loss = loss.sum()
+        elif self.reduction == 'mean':
+            loss = loss.mean()
+
+        return loss

BioMod/gene_perturbation/train_module.py

@@ -0,0 +1,325 @@
+"""
+================================================================================
+Genetic perturbation task training script (train_module.py)
+================================================================================
+
+[Task Overview]
+This script is used to train a model that predicts the effects of genetic perturbation. Gene perturbation refers to knockout or
+Overexpress a gene and observe its effect on the cell transcriptome. The goal of the model is: given a combination of genetic perturbations,
+Predict the gene expression profile of cells after perturbation.
+
+[Core process]
+1. Load perturbation data sets (such as Norman, Adamson and other public data sets)
+2. Initialize the pre-trained single-cell basic model (PerformerEncoderOnly_Module)
+3. Fine-tune the model on perturbed data and learn to predict the effect of perturbation
+4. Evaluate model performance (MSE, Pearson correlation coefficient and other indicators)
+
+[Key technical points]
+- Use DDP (DistributedDataParallel) for multi-GPU distributed training
+- Use module-based gene representation to group genes and then encode them
+- Supports multiple gene expression and gene2vec fusion methods
+
+[Adaptation instructions]
+This script is based on GEARS/train_dual_ddp_reduce_dec.py and adapted to the module-based model.
+Use a completely separate native module, avoiding any external gears_reduce imports.
+================================================================================
+"""
+
+import os
+# [Key] Disable Numba JIT compilation to avoid compilation conflicts in multi-process environments
+# Numba may cause problems during distributed training. Use pure Python to implement it after disabling it.
+os.environ['NUMBA_DISABLE_JIT'] = '1'
+import sys
+import time
+import argparse
+import pandas as pd
+import numpy as np
+import scanpy as sc
+from os.path import join as pjoin
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+import tempfile
+
+# ============================================================
+# [Key] Correctly set the order of sys.path
+# Local modules must take precedence and avoid importing into external utils.py
+# This is a common Python module import trap: if a module with the same name exists in the project root directory,
+# can be imported incorrectly, causing problems that are difficult to debug
+# ============================================================
+current_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# First add the project root directory (for performer_pytorch_cont)
+if project_root not in sys.path:
+    sys.path.append(project_root)
+
+# Finally add the current directory (make sure local modules take precedence)
+# Use append instead of insert(0), but since the current directory is empty there will be no conflict
+if current_dir not in sys.path:
+    sys.path.insert(0, current_dir)
+
+# ============================================================
+# [Key] Create an independent Numba cache directory for each process
+# In distributed training, multiple processes sharing the same cache directory can cause file conflicts.
+# This is a common problem with Numba in a DDP environment. An independent cache must be allocated for each rank.
+# ============================================================
+if "RANK" in os.environ:  # Distributed environment
+    cache_dir = f"/tmp/numba_cache_{os.environ['RANK']}"
+else:
+    cache_dir = tempfile.mkdtemp(prefix="numba_cache_")
+os.environ["NUMBA_CACHE_DIR"] = cache_dir
+
+# ============================================================
+# [Module Import] Use completely independent local modules
+# These local modules are to avoid conflicts with modules of the same name in other parts of the project
+# ============================================================
+from gears_reduce_local import GEARS
+from utils import print_sys
+from pertdata_local import PertData  # Local version, does not rely on external gears_reduce
+
+# Import module model
+from model_module import GEARS_Model_Module
+
+def main(parser):
+    """
+    main training function
+
+    [Training process]
+    1. Initialize the distributed training environment (DDP)
+    2. Loading and preprocessing perturbation data
+    3. Prepare data partitioning (training/validation/testing)
+    4. Initialize model and optimizer
+    5. Training loop + validation + testing
+    """
+    # Turn on anomaly detection to help debug gradient calculation problems
+    torch.autograd.set_detect_anomaly(True)
+
+    args = parser.parse_args()
+
+    # ============================================================
+    # [Distributed training initialization]
+    # DDP uses NCCL backend for inter-GPU communication
+    # local_rank: the GPU number of the current process on this machine
+    # RANK: global process number (different machines have different ranks during multi-machine training)
+    # ============================================================
+    local_rank = args.local_rank
+    rank = int(os.environ["RANK"])
+    is_master = (rank == 0)  # Rank 0 is the main process, responsible for printing logs and saving models
+
+    dist.init_process_group(backend='nccl', init_method='env://')
+    print('DDP init!')
+    local_rank = int(os.environ['LOCAL_RANK'])
+    device = f'cuda:{local_rank}'
+    torch.cuda.set_device(device)
+    world_size = torch.distributed.get_world_size()  # Total number of processes
+
+    print("local_rank: ", local_rank)
+
+    valid_every = args.valid_every  # Verify every epoch
+
+    # ============================================================
+    # [Data loading]
+    # The PertData class is responsible for loading and processing perturbation data
+    # Supported data sets: norman, adamson, dixit and other public data sets
+    # ============================================================
+    pert_data = PertData(args.data_dir)
+    # load dataset in paper: norman, adamson, dixit.
+    try:
+        if args.data_name in ['norman', 'adamson', 'dixit']:
+            pert_data.load(data_name=args.data_name)
+        else:
+            print('load data')
+            pert_data.load(data_path=pjoin(args.data_dir, args.data_name))
+    except:
+        # If the data format is not standard, try loading from h5ad file
+        adata = sc.read_h5ad(pjoin(args.data_dir, args.data_name+'.h5ad'))
+        adata.uns['log1p'] = {}
+        adata.uns['log1p']['base'] = None
+        pert_data.new_data_process(dataset_name=args.data_name, adata=adata)
+
+    # ============================================================
+    # [Gene Mapping]
+    # Check the overlap between the genes of the input data and the genes of the pre-trained model
+    # An overlap ratio that is too low may result in poor model performance.
+    # ============================================================
+    input_genes_ens_ids = pert_data.adata.var.index.tolist()
+    with open(args.scfm_genes_list_path, 'r') as f:
+        scfm_genes_ens_ids = [line.rstrip('\n') for line in f.readlines()]
+    if is_master:
+        print("target and scfm genes intersect: ", len(np.intersect1d(input_genes_ens_ids, scfm_genes_ens_ids)))
+
+    # ============================================================
+    # [Data division]
+    # The split parameter determines the data partitioning strategy:
+    # - 'simulation': simulates the division of real application scenarios, and the test set contains unseen gene combinations
+    # - 'combo_seen0/1/2': divided according to the number of genes seen in the combination
+    # train_gene_set_size: training set gene proportion
+    # ============================================================
+    pert_data.prepare_split(split=args.split, seed=args.seed, train_gene_set_size=args.train_gene_set_size)
+    # get dataloader with batch size
+    pert_data.get_dataloader(batch_size=args.batch_size, test_batch_size=args.test_batch_size)
+
+    # set up and train a model
+    if is_master:
+        print_sys("EXPERIMENT: " + args.run_name)
+
+    # ============================================================
+    # [Model initialization]
+    # The GEARS class encapsulates the complete process of training, verification, and testing
+    # ============================================================
+    gears_model = GEARS(pert_data,
+                        local_rank=local_rank,
+                        is_master=is_master,
+                        world_size=world_size,
+                        weight_bias_track=args.wandb,
+                        train_bs=args.batch_size,
+                        test_bs=args.test_batch_size,
+                        device=device,
+                        exp_name=args.run_name)
+
+    # ============================================================
+    # [Model parameter configuration]
+    # Use GEARS_Model_Module instead of GEARS_Model_Acc
+    # Key parameter description:
+    # - hidden_size: hidden layer dimension, affecting model capacity
+    # - finetune_method: fine-tuning method ('frozen' freezes pre-training parameters, None means fine-tuning all parameters)
+    # - accumulation_steps: the number of gradient accumulation steps, used to simulate larger batch size
+    # ============================================================
+    gears_model.model_initialize(
+        hidden_size=args.hidden_size,
+        model_type=args.model_type,
+        bin_set=args.bin_set,
+        load_path=args.singlecell_model_path,
+        finetune_method=args.finetune_method,
+        accumulation_steps=args.accumulation_steps,
+        mode=args.mode,
+        input_genes_ens_ids=input_genes_ens_ids,
+        scfm_genes_ens_ids=scfm_genes_ens_ids,
+        scfm_hyper_params_path=args.scfm_hyper_params_path,
+        scfm_ckpt_path=args.scfm_ckpt_path,
+        scfm_class=None,  # Do not use scfm_class, use module model directly
+        key_enc=None,
+        scfm_gene2vec_file=args.scfm_gene2vec_file,
+        record_pred=args.record_pred,
+        model_class=GEARS_Model_Module,  # Use module model class
+        # ============================================================
+        # [Module related parameters] The following parameters affect how the model processes module-based data
+        # module_data_path/module_data_dir: module divides the data path
+        # gene_mapping_path: mapping file from gene symbol to Ensembl ID
+        # use_module_level: whether to use module level features directly (default False, reverts to gene level)
+        # fusion_method: Read from pkl to ensure consistency with pre-training
+        # shuffle_module_genes: [Ablation experiment] Disrupt the gene distribution within the module
+        # ============================================================
+        module_data_path=args.module_data_path if hasattr(args, 'module_data_path') else None,
+        module_data_dir=args.module_data_dir if hasattr(args, 'module_data_dir') else None,
+        gene_mapping_path=args.gene_mapping_path if hasattr(args, 'gene_mapping_path') else '../gene2ens.csv',
+        scfm_genes_list_path=args.scfm_genes_list_path,
+        use_module_level=args.use_module_level if hasattr(args, 'use_module_level') else False,
+        shuffle_module_genes=args.shuffle_module_genes if hasattr(args, 'shuffle_module_genes') else False,
+        shuffle_seed=args.shuffle_seed if hasattr(args, 'shuffle_seed') else 2021,
+    )
+
+    gears_model.train(
+        epochs=args.epochs,
+        result_dir=args.result_dir,
+        lr=args.lr,
+        valid_every=valid_every
+    )
+
+    # save params
+    if is_master:
+        param_pd = pd.DataFrame(vars(args), index=['params']).T
+        param_pd.to_csv(f'{args.result_dir}/params.csv')
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GEARS with PerformerEncoderOnly_Module')
+
+    # ============================================================
+    # [Distributed training parameters]
+    # local-rank: automatically set by torch.distributed.launch, no need to specify manually
+    # ============================================================
+    parser.add_argument("--local-rank", type=int, default=-1, help='Local process rank.')
+
+    # ============================================================
+    # [Data related parameters]
+    # data_dir: data storage directory
+    # data_name: Data set name (norman/adamson/dixit is a public data set)
+    # split: data partitioning strategy, simulation is the most commonly used
+    # train_gene_set_size: The proportion of genes in the training set, which affects the evaluation of model generalization ability
+    # ============================================================
+    parser.add_argument('--data_dir', type=str, default='../../GEARS/data')
+    parser.add_argument('--data_name', type=str, default='norman')
+    parser.add_argument('--split', type=str, default='simulation')
+    parser.add_argument('--result_dir', type=str, default='./results')
+    parser.add_argument('--seed', type=int, default=1)
+
+    # ============================================================
+    # [Training hyperparameters]
+    # epochs: number of training rounds, usually 15-20 rounds is enough
+    # batch_size: Batch size, adjusted according to GPU memory
+    # lr: learning rate, usually a smaller learning rate (1e-3 or less) is used during fine-tuning
+    # accumulation_steps: Number of gradient accumulation steps, used to simulate larger batch size
+    # ============================================================
+    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--batch_size', type=int, default=32)
+    parser.add_argument('--valid_every', type=int, default=1)
+    parser.add_argument('--test_batch_size', type=int, default=128)
+    parser.add_argument('--train_gene_set_size', type=float, default=0.75)
+    parser.add_argument('--hidden_size', type=int, default=64)
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--model_type', type=str, default=None)
+    parser.add_argument('--bin_set', type=str, default=None)
+    parser.add_argument('--singlecell_model_path', type=str, default=None)
+    parser.add_argument('--finetune_method', type=str, default=None)
+    parser.add_argument('--mode', type=str, default='v1')
+    parser.add_argument('--accumulation_steps', type=int, default=1)
+    parser.add_argument('-wandb', action='store_true', help='Enable Weights & Biases logging')
+    parser.add_argument('-record_pred', action='store_true', help='Record predictions for analysis')
+
+    # ============================================================
+    # [Related parameters of pre-training model]
+    # scfm_genes_list_path: list of genes used by the pre-trained model
+    # scfm_hyper_params_path: Hyperparameter configuration file of the pre-trained model
+    # scfm_ckpt_path: checkpoint file of the pre-trained model
+    # scfm_gene2vec_file: gene2vec pre-trained embedding file
+    # ============================================================
+    parser.add_argument('--scfm_genes_list_path', type=str,
+                        default="../../selected_genes_27k.txt")
+    parser.add_argument('--scfm_hyper_params_path', type=str,
+                        default="../ckpts_performer_module_test/performer_module_test.pkl")
+    parser.add_argument('--scfm_ckpt_path', type=str,
+                        default="../ckpts_performer_module_test/performer_module_test_epoch6.pth")
+    parser.add_argument('--run_name', type=str,
+                        default="default_run_module")
+    parser.add_argument('--scfm_gene2vec_file', type=str,
+                        default='../../selected_gene2vec_27k.npy')
+
+    # ============================================================
+    # [Module related parameters] These parameters affect how the model handles module-based gene representation.
+    # module_data_path: single module JSON file path
+    # module_data_dir: module JSON file directory
+    # gene_mapping_path: mapping file from gene symbol to Ensembl ID
+    # use_module_level: whether to use module level features directly (experimental feature)
+    # fusion_method: Read from pkl to ensure consistency with pre-training
+    # shuffle_module_genes: [Ablation experiment] Disrupt the genes in the module to verify the importance of the module structure
+    # ============================================================
+    parser.add_argument('--module_data_path', type=str, default=None,
+                        help='Path to module JSON file (single file)')
+    parser.add_argument('--module_data_dir', type=str, default='../module_data',
+                        help='Path to directory containing module JSON files')
+    parser.add_argument('--gene_mapping_path', type=str, default='../../gene2ens.csv',
+                        help='Path to gene symbol to Ensembl ID mapping CSV')
+    parser.add_argument('--use_module_level', action='store_true',
+                        help='Use module-level features directly instead of recovering to gene level. '
+                             'Note: This requires compatible downstream architecture (not fully implemented).')
+    parser.add_argument('--shuffle_module_genes', action='store_true',
+                        help='[Ablation experiment] Additionally shuffle module_sets after restoring them from the pretrained PKL. '
+                             'Do not pass this parameter when reproducing the experiment normally.')
+    parser.add_argument('--shuffle_seed', type=int, default=2021,
+                        help='Only takes effect when --shuffle_module_genes is turned on and is used for additional downstream shuffling.')
+    parser.add_argument('--lr', type=float, default=1e-3)
+
+    main(parser)

BioMod/gene_perturbation/utils.py

@@ -0,0 +1,489 @@
+"""
+================================================================================
+Utility function module (utils.py)
+================================================================================
+
+[Module function]
+Provides various tool functions required for genetic perturbation tasks, including:
+
+1. Perturbation analysis function: analyze the gene name of single gene/combination perturbation
+2. Gene relationship graph construction: co-expression graph and GO similarity graph
+3. Loss function: training loss and uncertainty loss
+4. Genetic interaction (GI) analysis: Calculate genetic interaction coefficients
+
+[Key concepts]
+- Perturbation format: "gene1+gene2", such as "ctrl+MAPK1" or "MAPK1+KRAS"
+- Co-expression map: constructed based on gene expression correlations in the data
+- GO graph: Jaccard similarity construction based on Gene Ontology annotation
+
+[Key parameters that may affect the results]
+- coexpress_threshold: Threshold for co-expression graph construction
+- dict_filter: training loss is only calculated on non-zero genes
+- direction_lambda: direction consistency loss weight
+================================================================================
+"""
+
+import torch
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import TheilSenRegressor
+import torch.nn as nn
+
+import networkx as nx
+from tqdm import tqdm
+import pickle
+import sys, os
+import requests
+from torch_geometric.data import Data
+
+from zipfile import ZipFile
+
+import statsmodels.api as sm
+from sklearn.linear_model import LinearRegression, TheilSenRegressor
+# ============================================================
+# Lazy loading of dcor to avoid numba compatibility issues in DDP multi-process environments
+# dcor will be imported when needed (in the get_coeffs function)
+# ============================================================
+from sklearn.metrics import r2_score
+
+def parse_single_pert(i):
+    """
+    Parsing single gene perturbation strings
+
+    [Format]
+    - "ctrl+GENE": The perturbed gene is GENE
+    - "GENE+ctrl": The perturbed gene is GENE
+
+    [return]
+    Perturbed gene name
+    """
+    a = i.split('+')[0]
+    b = i.split('+')[1]
+    if a == 'ctrl':
+        pert = b
+    else:
+        pert = a
+    return pert
+
+def filter_pert_cond_in_go(condition, exclude_pert):
+    """
+    Filter perturbations not in the GO graph
+
+    [use]
+    Genes without annotations in the GO diagram cannot predict their perturbation effects.
+    These disturbing conditions need to be eliminated during data loading.
+    """
+    if condition == 'ctrl':
+        return True
+    else:
+        if condition in exclude_pert:
+            return False
+        else:
+            return True
+
+def parse_combo_pert(i):
+    """
+    Parsing combined perturbation strings
+
+    [Format]
+    - "GENE1+GENE2": perturb two genes simultaneously
+
+    [return]
+    (gene1, gene2) tuple
+    """
+    return i.split('+')[0], i.split('+')[1]
+
+def combine_res(res_1, res_2):
+    """Merge two evaluation result dictionaries"""
+    res_out = {}
+    for key in res_1:
+        res_out[key] = np.concatenate([res_1[key], res_2[key]])
+    return res_out
+
+def parse_any_pert(p):
+    """
+    Parse arbitrary perturbation strings (single genes or combinations)
+
+    [return]
+    Gene list:
+    - Single gene perturbation: [gene]
+    - Combinatorial perturbation: [gene1, gene2]
+    """
+    if ('ctrl' in p) and (p != 'ctrl'):
+        return [parse_single_pert(p)]
+    elif 'ctrl' not in p:
+        out = parse_combo_pert(p)
+        return [out[0], out[1]]
+
+def np_pearson_cor(x, y):
+    """
+    Calculate Pearson correlation coefficient matrix
+
+    [use]
+    Used to construct co-expressed gene relationship diagrams
+    """
+    xv = x - x.mean(axis=0)
+    yv = y - y.mean(axis=0)
+    xvss = (xv * xv).sum(axis=0)
+    yvss = (yv * yv).sum(axis=0)
+    result = np.matmul(xv.transpose(), yv) / np.sqrt(np.outer(xvss, yvss))
+    return np.maximum(np.minimum(result, 1.0), -1.0)
+
+
+def dataverse_download(url, save_path):
+    """dataverse download helper with progress bar
+
+    Args:
+        url (str): the url of the dataset
+        path (str): the path to save the dataset
+    """
+    
+    if os.path.exists(save_path):
+        print_sys('Found local copy...')
+    else:
+        print_sys("Downloading...")
+        response = requests.get(url, stream=True)
+        total_size_in_bytes= int(response.headers.get('content-length', 0))
+        block_size = 1024
+        progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+        with open(save_path, 'wb') as file:
+            for data in response.iter_content(block_size):
+                progress_bar.update(len(data))
+                file.write(data)
+        progress_bar.close()
+
+        
+def zip_data_download_wrapper(url, save_path, data_path):
+
+    if os.path.exists(save_path):
+        print_sys('Found local copy...')
+    else:
+        dataverse_download(url, save_path + '.zip')
+        print_sys('Extracting zip file...')
+        with ZipFile((save_path + '.zip'), 'r') as zip:
+            zip.extractall(path = data_path)
+        print_sys("Done!")  
+        
+def get_go_auto(gene_list, data_path, data_name):
+    go_path = os.path.join(data_path, data_name, 'go.csv')
+    
+    if os.path.exists(go_path):
+        return pd.read_csv(go_path)
+    else:
+        with open(os.path.join('/home/ding.bai/pert_new/GEARS/data_all/gene2go_all.pkl'), 'rb') as f:
+            gene2go = pickle.load(f)
+
+        gene2go = {i: list(gene2go[i]) for i in gene_list if i in gene2go}
+        edge_list = []
+        for g1 in tqdm(gene2go.keys()):
+            for g2 in gene2go.keys():
+                edge_list.append((g1, g2, len(np.intersect1d(gene2go[g1], gene2go[g2]))/len(np.union1d(gene2go[g1], gene2go[g2]))))
+
+        edge_list_filter = [i for i in edge_list if i[2] > 0]
+        further_filter = [i for i in edge_list if i[2] > 0.1]
+        df_edge_list = pd.DataFrame(further_filter).rename(columns = {0: 'gene1', 1: 'gene2', 2: 'score'})
+
+        df_edge_list = df_edge_list.rename(columns = {'gene1': 'source', 'gene2': 'target', 'score': 'importance'})
+
+        df_edge_list.to_csv(go_path, index = False)        
+        return df_edge_list
+
+def get_go(df_gene2go):
+    df_gene2go['Entry name'] = df_gene2go['Entry name'].apply(lambda x: x.split('_')[0])
+    df_gene2go = df_gene2go[df_gene2go['Gene ontology IDs'].notnull()]
+    df_gene2go = df_gene2go.rename(columns = {[i for i in df_gene2go.columns.values if 'yourlist' in i][0]: 'gene_id'})
+
+    geneid2go = dict(df_gene2go[['gene_id', 'Gene ontology IDs']].values)
+
+    gene2go = {}
+    for i,j in geneid2go.items():
+        j = [k.strip() for k in j.split(';')]
+        for k in i.split(','):
+            gene2go[ensembl2genename[k]] = j
+
+    from tqdm import tqdm
+    edge_list = []
+    for g1 in tqdm(gene2go.keys()):
+        for g2 in gene2go.keys():
+
+            edge_list.append((g1, g2, len(np.intersect1d(gene2go[g1], gene2go[g2]))/len(np.union1d(gene2go[g1], gene2go[g2]))))
+
+    edge_list_filter = [i for i in edge_list if i[2] > 0]
+    further_filter = [i for i in edge_list if i[2] > 0.1]
+    df_edge_list = pd.DataFrame(further_filter).rename(columns = {0: 'gene1', 1: 'gene2', 2: 'score'})
+
+    df_edge_list = df_edge_list.rename(columns = {'gene1': 'source', 'gene2': 'target', 'score': 'importance'})
+
+    return df_edge_list
+
+class GeneSimNetwork():
+    '''
+    Represents a gene similarity network. Has the following attributes:
+        - self.edge_list
+        - self.gene_list
+        - self.G
+        - self.edge_index (Tensor, can be used to set G_go in model_initialize)
+        - self.edge_weight (Tensor, can be used to set G_go_weight in model_initialize)
+    '''
+    def __init__(self, edge_list, gene_list, node_map):
+        # Keeps edges if both source and target exist in node_map
+        keep_edges = np.zeros(len(edge_list)).astype(bool)
+        for i, (source, target, importance) in edge_list.iterrows():
+            keep_edges[i] = source in node_map and target in node_map
+        self.edge_list = edge_list[keep_edges]
+
+        # Constructs NetworkX graph
+        self.G = nx.from_pandas_edgelist(self.edge_list, source='source',
+                        target='target', edge_attr=['importance'],
+                        create_using=nx.DiGraph())
+        self.gene_list = gene_list
+
+        # Add any isolated nodes (represented in the tensor with self edges)
+        for n in self.gene_list:
+            added_nodes = []
+            if n not in self.G.nodes():
+                self.G.add_node(n)
+                added_nodes.append(n)
+        if len(added_nodes) > 0:
+            print('The following nodes were not present in the edge list and have been added to the graph as self-edges:', added_nodes)
+        
+        # Convert graph edges to tensor
+        edge_index_ = [(node_map[e[0]], node_map[e[1]]) for e in self.G.edges]
+        self.edge_index = torch.tensor(edge_index_, dtype=torch.long).T
+        
+        # Convert edge weights to tensor
+        edge_attr = nx.get_edge_attributes(self.G, 'importance') 
+        importance = np.array([edge_attr[e] for e in self.G.edges])
+        self.edge_weight = torch.Tensor(importance)
+
+
+def get_similarity_network(network_type, adata, threshold, k, gene_list, data_path, data_name, split, seed, train_gene_set_size, set2conditions):
+    
+    if network_type == 'co-express':
+        df_out = get_coexpression_network_from_train(adata, threshold, k, data_path, data_name, split, seed, train_gene_set_size, set2conditions)
+    elif network_type == 'go':
+        df_jaccard = get_go_auto(gene_list, data_path, data_name)
+        df_out = df_jaccard.groupby('target').apply(lambda x: x.nlargest(k + 1,['importance'])).reset_index(drop = True)
+
+    return df_out
+
+def get_coexpression_network_from_train(adata, threshold, k, data_path, data_name, split, seed, train_gene_set_size, set2conditions):
+    
+    fname = os.path.join(os.path.join(data_path, data_name), split + '_' + str(seed) + '_' + str(train_gene_set_size) + '_' + str(threshold) + '_' + str(k) + '_co_expression_network.csv')
+
+    
+    if os.path.exists(fname):
+        return pd.read_csv(fname)
+    else:
+        gene_list = [f for f in adata.var.gene_name.values]
+        idx2gene = dict(zip(range(len(gene_list)), gene_list)) 
+        X = adata.X
+        train_perts = set2conditions['train']
+        X_tr = X[np.isin(adata.obs.condition, [i for i in train_perts if 'ctrl' in i])]
+        gene_list = adata.var['gene_name'].values
+
+        X_tr = X_tr.toarray()
+        out = np_pearson_cor(X_tr, X_tr)
+        out[np.isnan(out)] = 0
+        out = np.abs(out)
+
+        out_sort_idx = np.argsort(out)[:, -(k + 1):]
+        out_sort_val = np.sort(out)[:, -(k + 1):]
+
+        df_g = []
+        for i in range(out_sort_idx.shape[0]):
+            target = idx2gene[i]
+            for j in range(out_sort_idx.shape[1]):
+                df_g.append((idx2gene[out_sort_idx[i, j]], target, out_sort_val[i, j]))
+
+        df_g = [i for i in df_g if i[2] > threshold]
+
+        df_co_expression = pd.DataFrame(df_g).rename(columns = {0: 'source', 1: 'target', 2: 'importance'})
+        df_co_expression.to_csv(fname, index = False)
+        return df_co_expression
+    
+
+def uncertainty_loss_fct(pred, logvar, y, perts, reg = 0.1, ctrl = None, direction_lambda = 1e-3, dict_filter = None):
+    """
+    Uncertain version loss:
+    1) The error term is weighted by logvar (heteroscedastic modeling)
+    2) Optional direction penalty (relative to the up/down direction of control)
+    3) `dict_filter` is only calculated on a non-zero gene subset, and the caliber is consistent with the evaluation
+    """
+    gamma = 2                     
+    perts = np.array(perts)
+    losses = torch.tensor(0.0, requires_grad=True).to(pred.device)
+    for p in set(perts):
+        if p!= 'ctrl':
+            # Training only on non-zero genes of this perturbation can reduce the impact of dropout noise.
+            retain_idx = dict_filter[p]
+            pred_p = pred[np.where(perts==p)[0]][:, retain_idx]
+            y_p = y[np.where(perts==p)[0]][:, retain_idx]
+            logvar_p = logvar[np.where(perts==p)[0]][:, retain_idx]
+        else:
+            pred_p = pred[np.where(perts==p)[0]]
+            y_p = y[np.where(perts==p)[0]]
+            logvar_p = logvar[np.where(perts==p)[0]]
+                         
+        # uncertainty based loss
+
+        losses += torch.sum((pred_p - y_p)**(2 + gamma) + reg * torch.exp(-logvar_p) * (pred_p - y_p)**(2 + gamma))/pred_p.shape[0]/pred_p.shape[1]
+                         
+        # direction loss                 
+        if p!= 'ctrl':
+            # Directional consistency term: The constraint prediction is consistent with the true direction of change relative to control.
+            losses += torch.sum(direction_lambda * (torch.sign(y_p - ctrl[retain_idx]) - torch.sign(pred_p - ctrl[retain_idx]))**2)/pred_p.shape[0]/pred_p.shape[1]
+        else:
+            losses += torch.sum(direction_lambda * (torch.sign(y_p - ctrl) - torch.sign(pred_p - ctrl))**2)/pred_p.shape[0]/pred_p.shape[1]
+            
+    return losses/(len(set(perts)))
+
+def loss_fct(pred, y, perts, ctrl = None, direction_lambda = 1e-3, dict_filter = None):
+    """
+    Main training loss (without uncertainty):
+    - Basic item: MSE
+    - Optional direction item: Penalize upward/downward direction errors
+    - `dict_filter` will change the training gene set and is an important switch that affects the results.
+    """
+    gamma = 2
+    mse_p = torch.nn.MSELoss()
+    perts = np.array(perts)
+    losses = torch.tensor(0.0).to(pred.device)
+
+    for p in set(perts):
+        pert_idx = np.where(perts == p)[0]
+        
+
+        # during training, we remove the all zero genes into calculation of loss. this gives a cleaner direction loss. empirically, the performance stays the same.
+
+        # during training, we remove the all zero genes into calculation of loss.
+        # this gives a cleaner direction loss. empirically, the performance stays the same.
+
+        if p!= 'ctrl':
+            # Generally more stable when consistent with the evaluated non-zero / non-dropout settings.
+            retain_idx = dict_filter[p]
+            pred_p = pred[pert_idx][:, retain_idx]
+            y_p = y[pert_idx][:, retain_idx]
+        else:
+            pred_p = pred[pert_idx]
+            y_p = y[pert_idx]
+
+        
+        losses += torch.sum((pred_p - y_p)**(2 + gamma))/pred_p.shape[0]/pred_p.shape[1]
+                         
+        ## direction loss
+        if (p!= 'ctrl'):
+            losses += torch.sum(direction_lambda * (torch.sign(y_p - ctrl[retain_idx]) - torch.sign(pred_p - ctrl[retain_idx]))**2)/pred_p.shape[0]/pred_p.shape[1]
+        else:
+            losses += torch.sum(direction_lambda * (torch.sign(y_p - ctrl) - torch.sign(pred_p - ctrl))**2)/pred_p.shape[0]/pred_p.shape[1]
+
+    return losses/(len(set(perts)))
+
+
+def print_sys(s):
+    """system print
+
+    Args:
+        s (str): the string to print
+    """
+    print(s, flush = True, file = sys.stderr)
+    
+def create_cell_graph_for_prediction(X, pert_idx, pert_gene):
+
+
+    # If perturbations will be represented as node features
+    pert_feats = np.zeros(len(X))
+    for p in pert_idx:
+        pert_feats[int(np.abs(p))] = np.sign(p)
+    feature_mat = torch.Tensor(np.vstack([X, pert_feats])).T
+
+    return Data(x=feature_mat, pert=pert_gene)
+    
+
+def create_cell_graph_dataset_for_prediction(pert_gene, ctrl_adata, gene_names, device, num_samples = 300):
+    Xs = []
+
+    # Get the indices (and signs) of applied perturbation
+    pert_idx = [np.where(p == np.array(gene_names))[0][0] for p in pert_gene]
+
+    Xs = ctrl_adata[np.random.randint(0, len(ctrl_adata), num_samples), :].X.toarray()
+    # Create cell graphs
+    cell_graphs = [create_cell_graph_for_prediction(X, pert_idx, pert_gene).to(device) for X in Xs]
+    return cell_graphs
+
+##
+##GI related utils
+##
+
+def get_coeffs(singles_expr, first_expr, second_expr, double_expr):
+    # ============================================================
+    # Lazy loading of dcor to avoid numba conflicts when importing into multiple processes
+    # Import dcor only when the function is called
+    # ============================================================
+    try:
+        from dcor import distance_correlation
+    except Exception as e:
+        # Use alternative if dcor import fails (numba compatibility issue)
+        import warnings
+        warnings.warn(f"dcor import failed: {e}. Using fallback distance correlation implementation.")
+        # Use a simplified scipy/numpy based implementation as an alternative
+        def distance_correlation(x, y):
+            """Simplified version of distance correlation coefficient, does not rely on numba"""
+            from scipy.spatial.distance import pdist, squareform
+            x = np.asarray(x)
+            y = np.asarray(y)
+            a = squareform(pdist(x))
+            b = squareform(pdist(y))
+            A = a - a.mean(axis=0)[None, :] - a.mean(axis=1)[:, None] + a.mean()
+            B = b - b.mean(axis=0)[None, :] - b.mean(axis=1)[:, None] + b.mean()
+            dcov2 = (A * B).sum() / (A.size * A.size)
+            dvar_x = (A * A).sum() / (A.size * A.size)
+            dvar_y = (B * B).sum() / (B.size * B.size)
+            if dvar_x == 0 or dvar_y == 0:
+                return 0.0
+            return np.sqrt(dcov2 / np.sqrt(dvar_x * dvar_y))
+
+    results = {}
+    results['ts'] = TheilSenRegressor(fit_intercept=False,
+                          max_subpopulation=1e5,
+                          max_iter=1000,
+                          random_state=1000)
+    X = singles_expr
+    y = double_expr
+    results['ts'].fit(X, y.ravel())
+    Zts = results['ts'].predict(X)
+    results['c1'] = results['ts'].coef_[0]
+    results['c2'] = results['ts'].coef_[1]
+    results['mag'] = np.sqrt((results['c1']**2 + results['c2']**2))
+
+    results['dcor'] = distance_correlation(singles_expr, double_expr)
+    results['dcor_singles'] = distance_correlation(first_expr, second_expr)
+    results['dcor_first'] = distance_correlation(first_expr, double_expr)
+    results['dcor_second'] = distance_correlation(second_expr, double_expr)
+    results['corr_fit'] = np.corrcoef(Zts.flatten(), double_expr.flatten())[0,1]
+    results['dominance'] = np.abs(np.log10(results['c1']/results['c2']))
+    results['eq_contr'] = np.min([results['dcor_first'], results['dcor_second']])/\
+                        np.max([results['dcor_first'], results['dcor_second']])
+
+    return results
+
+def get_GI_params(preds, combo):
+    
+    singles_expr = np.array([preds[combo[0]], preds[combo[1]]]).T
+    first_expr = np.array(preds[combo[0]]).T
+    second_expr = np.array(preds[combo[1]]).T
+    double_expr = np.array(preds[combo[0]+'_'+combo[1]]).T
+    
+    return get_coeffs(singles_expr, first_expr, second_expr, double_expr)
+
+def get_GI_genes_idx(adata, GI_gene_file):
+    # Genes used for linear model fitting
+    GI_genes = np.load(GI_gene_file, allow_pickle=True)
+    GI_genes_idx = np.where([g in GI_genes for g in adata.var.gene_name.values])[0]
+    
+    return GI_genes_idx
+
+def get_mean_control(adata):
+    cols = adata.var.gene_name.values.astype('str')
+    mean_ctrl_exp = adata[adata.obs['condition'] == 'ctrl'].to_df().mean()
+    return mean_ctrl_exp

BioMod/interpretability/README.md

@@ -0,0 +1,412 @@
+# Module Interpretability Experiments for scLong (performer_module)
+
+This module contains interpretability experiments designed specifically for the
+**module level** of scLong. The goal is to demonstrate, using structural,
+semantic, and functional evidence, that module design has a paradigm-level
+advantage over gene-level foundation models such as GenePT, scGPT, and
+scFoundation.
+
+## Experiment Roadmap, Ordered by Priority
+
+| # | Script | Core question | Status |
+| --- | --- | --- | --- |
+| 1 | `module_geometry.py` | Does the module latent space self-organize by pathway family,  | Completed |
+| 2 | `module_text_alignment/module_text_alignment.py` | Do module embeddings align with natural-language pathway descriptions,  | Implemented |
+| 3 | `module_importance/module_importance.py` | Does task-driven module importance match known pathway-cell-type or pathway-disease associations,  | Implemented |
+| 4 | `module_vs_gene.py` | Is scLong better than GenePT/scGPT for pathway-level interpretability,  | TODO |
+| 6 | `attention/module_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 (`module_geometry.py`)
+
+**Hypothesis:** the scLong encoder is trained only 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_modules.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. scLong 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
+│   ├── module_crosstalk.py
+│   ├── module_gene_attention.py
+│   ├── README.md
+│   └── README_gene_attention.md
+├── module_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
+│   ├── module_importance.py
+│   ├── module_importance_publication.py
+│   ├── pub_utils.py
+│   ├── random_baseline.py
+│   ├── README.md
+│   ├── RESULTS.md
+│   ├── tcell_reactome_analysis.py
+│   └── TODO.md
+├── module_text_alignment/                               # Semantic alignment of module representations and pathway text
+│   ├── __init__.py
+│   ├── DESIGN.md
+│   ├── fetch_reactome_descriptions.py
+│   ├── MANUSCRIPT_SECTION.md
+│   ├── module_text_alignment.py
+│   ├── module_text_alignment_extended.py
+│   ├── module_text_alignment_publication.py
+│   └── reactome_descriptions.json
+├── __init__.py
+├── module_geometry.py
+├── README.md
+└── utils.py
+```
+
+## Command-Line Usage
+
+Run scripts from `BioMod/interpretability/` within the project root; relative paths below are
+based on that directory.
+
+```bash
+# General h5ad input: cell annotation, zero-shot, or any annotated dataset.
+python module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --gene2ens_file        ../../gene2ens.csv \
+    --data_path            ../../zero-shot-batch/pancreas_scib.h5ad \
+    --data_mode            h5ad \
+    --gene_name_column     auto \
+    --batch_size           640 \
+    --output_dir           ./results/geometry/pancreas_scib
+```
+
+For pretraining data, where each cell is an `.npz` file already in the 27,874
+model-gene space:
+
+```bash
+python module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --data_path            ../../../../share/cell_files \
+    --data_mode            pretrain_npz \
+    --cell_num             48024242 \
+    --max_cells            50000 \
+    --batch_size           512 \
+    --output_dir           ./results/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 module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch2.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --data_path            ../../../share/cell_files \
+    --data_mode            pretrain_npz \
+    --cell_num             49738024 \
+    --max_cells            100000 \
+    --batch_size           16 \
+    --output_dir           ./results/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_module_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_module_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 `--performer_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_module_genes` to the original command and write to a control output
+directory.
+
+For h5ad data, using `pancreas_scib` as an example:
+
+```bash
+python module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --gene2ens_file        ../../gene2ens.csv \
+    --data_path            ../../zero-shot-batch/pancreas_scib.h5ad \
+    --data_mode            h5ad \
+    --gene_name_column     auto \
+    --batch_size           640 \
+    --shuffle_module_genes \
+    --output_dir           ./results/geometry/pancreas_scib_shuffle
+```
+
+For pretraining npz data:
+
+```bash
+python module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --gene2ens_file        ../../gene2ens.csv \
+    --data_path            ../../../../share/cell_files \
+    --data_mode            pretrain_npz \
+    --cell_num             48024242 \
+    --max_cells            50000 \
+    --batch_size           640 \
+    --shuffle_module_genes \
+    --output_dir           ./results/geometry/pretrain_subset_shuffle
+```
+
+How `--shuffle_module_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 `module_geometry.py`
+with it.
+
+**Step 1: Generate `_init.pth`**
+
+Use the pretraining script with `--save_init_ckpt_and_exit`. The command
+initializes the model, saves a random-weight checkpoint and the matching config
+PKL, then exits immediately without entering training.
+
+```bash
+cd /mnt/c/Users/250010166/Desktop/single_cell/scLong/performer_module
+
+torchrun --nproc_per_node=1 performer_module.py \
+    --model_name   mae_tost_48M_cross_attn_only_pure_memory \
+    --ckpt_dir     ./mae_tost_48M_cross_attn_only_pure_memory \
+    --gene_list_path ../selected_genes_27k.txt \
+    --gene2vec_data_path ../selected_gene2vec_27k.npy \
+    --module_data_dir module_data_all \
+    --gene_mapping_path ../gene2ens.csv \
+    --data_path    ../../../share/cell_files \
+    --cell_num     48024242 \
+    --enc_depth    32 \
+    --enc_heads    16 \
+    --base_dim     1024 \
+    --pretrain_arch mae \
+    --decoder_type cross_attn \
+    --save_init_ckpt_and_exit
+```
+
+The following files are generated under
+`mae_tost_48M_cross_attn_only_pure_memory/`:
+
+- `mae_tost_48M_cross_attn_only_pure_memory_init.pth`, randomly initialized weights.
+- `mae_tost_48M_cross_attn_only_pure_memory_init.pkl`, model config.
+
+**Step 2: Run `module_geometry.py` with `_init.pth`**
+
+For h5ad data:
+
+```bash
+cd /mnt/c/Users/250010166/Desktop/single_cell/scLong/performer_module/interpretability
+
+python module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory_init/mae_tost_48M_cross_attn_only_pure_memory_init.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory_init/mae_tost_48M_cross_attn_only_pure_memory_init.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --gene2ens_file        ../../gene2ens.csv \
+    --data_path            ../../zero-shot-batch/pancreas_scib.h5ad \
+    --data_mode            h5ad \
+    --gene_name_column     auto \
+    --batch_size           640 \
+    --output_dir           ./results/geometry/pancreas_scib_init
+```
+
+For pretraining npz data:
+
+```bash
+python module_geometry.py \
+    --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory_init/mae_tost_48M_cross_attn_only_pure_memory_init.pth \
+    --config_pkl           ../mae_tost_48M_cross_attn_only_pure_memory_init/mae_tost_48M_cross_attn_only_pure_memory_init.pkl \
+    --module_data_dir       ../module_data_all \
+    --gene_list_path       ../../selected_genes_27k.txt \
+    --gene_mapping_path    ../../gene2ens.csv \
+    --gene2vec_data_path   ../../selected_gene2vec_27k.npy \
+    --gene2ens_file        ../../gene2ens.csv \
+    --data_path            ../../../../share/cell_files \
+    --data_mode            pretrain_npz \
+    --cell_num             48024242 \
+    --max_cells            50000 \
+    --batch_size           640 \
+    --output_dir           ./results/geometry/pretrain_subset_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 (`module_text_alignment/`)
+
+**Hypothesis:** the structure of the scLong 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_modules.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
+cd performer_module/interpretability/module_text_alignment
+
+# Download the sentence encoder.
+wget https://hf-mirror.com/hfd/hfd.sh && chmod a+x hfd.sh && export HF_ENDPOINT=https://hf-mirror.com
+./hfd.sh pritamdeka/S-PubMedBert-MS-MARCO --local-dir ./S-PubMedBert-MS-MARCO
+
+# Main experiment: RSA + CKA + LOO retrieval.
+python module_text_alignment.py \
+    --module_embedding_path ../results/geometry/pretrain_subset/module_embeddings.npy \
+    --module_json ../../module_data_all/merged_modules.json \
+    --sentence_model ./S-PubMedBert-MS-MARCO \
+    --n_permutations 1000 \
+    --output_dir ../results/text_alignment/pretrain_subset_v2
+
+# Extended analysis: broad Reactome categories and Hallmark deep dive.
+python module_text_alignment_extended.py \
+    --module_embedding_path ../results/geometry/pretrain_subset/module_embeddings.npy \
+    --module_json ../../module_data_all/merged_modules.json \
+    --sentence_model ./S-PubMedBert-MS-MARCO \
+    --descriptions_file ./reactome_descriptions.json \
+    --n_permutations 1000 \
+    --output_dir ../results/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

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

BioMod/interpretability/attention/README.md

@@ -0,0 +1,108 @@
+# Experiment 6: Module Crosstalk Graph vs Pathway Interaction Database
+
+This experiment tests whether module-module interactions learned by the scLong
+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
+
+```bash
+cd /path/to/performer_module/interpretability/attention
+
+# Full run, methods A + C.
+python module_crosstalk.py \
+    --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../zero-shot-batch/pancreas_scib.h5ad \
+    --data_mode h5ad \
+    --batch_size 640 \
+    --jacobian_cells 200 \
+    --output_dir ../results/crosstalk/pancreas_max \
+    --pool_mode max
+
+# Method A only, for quick validation.
+python module_crosstalk.py \
+    --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../zero-shot-batch/pancreas_scib.h5ad \
+    --data_mode h5ad \
+    --batch_size 64 \
+    --skip_jacobian \
+    --output_dir ../results/crosstalk/pancreas_cosine_only
+
+# Pretraining npz data.
+python module_crosstalk.py \
+    --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../../share/cell_files \
+    --data_mode pretrain_npz \
+    --cell_num 49738024 \
+    --max_cells 5000 \
+    --batch_size 16 \
+    --jacobian_cells 100 \
+    --output_dir ../results/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

@@ -0,0 +1,114 @@
+# Gene Attention Interpretability: ModuleEmbedding
+
+Analyze gene contributions within modules using the content-based attention
+pooling in ModuleEmbedding.
+
+## 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
+
+### Basic Usage on the Pancreas Dataset with Built-In Markers
+
+```bash
+cd ../data/250010166/scLong/performer_module/interpretability/attention
+
+python module_gene_attention.py \
+    --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../zero-shot-batch/pancreas_scib.h5ad \
+    --celltype_column celltype \
+    --batch_size 640 \
+    --top_k_genes 20 \
+    --output_dir ../results/gene_attention/pancreas
+```
+
+### With DEG Comparison
+
+```bash
+python module_gene_attention.py \
+    --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../zero-shot-batch/pancreas_scib.h5ad \
+    --celltype_column celltype \
+    --batch_size 64 \
+    --deg_csv /path/to/deg_results.csv \
+    --output_dir ../results/gene_attention/pancreas_with_deg
+```
+
+DEG CSV columns must include `gene`, `cell_type`, `log2fc`, and `pval`.
+
+### Custom Marker Database
+
+```bash
+python module_gene_attention.py \
+    ... \
+    --marker_db /path/to/markers.csv \
+    --output_dir ../results/gene_attention/custom_markers
+```
+
+Marker CSV columns must include `cell_type` and `gene_symbol`.
+
+### Limit Cell Count for Faster Testing
+
+```bash
+python module_gene_attention.py \
+    ... \
+    --max_cells 5000 \
+    --output_dir ../results/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_module_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/module_crosstalk.py

@@ -0,0 +1,605 @@
+"""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:
+  python module_crosstalk.py \
+      --performer_checkpoint ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+      --config_pkl ../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+      --module_data_dir ../module_data_all \
+      --gene_list_path ../../selected_genes_27k.txt \
+      --gene_mapping_path ../../gene2ens.csv \
+      --gene2vec_data_path ../../selected_gene2vec_27k.npy \
+      --data_path ../../zero-shot-batch/pancreas_scib.h5ad \
+      --data_mode h5ad \
+      --batch_size 16 \
+      --jacobian_cells 200 \
+      --output_dir ./results/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()
+PERFORMER_MODULE = HERE.parents[2]
+REPO_ROOT = HERE.parents[3]
+for p in (str(PERFORMER_MODULE), str(REPO_ROOT)):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+from cell_annotation.performer_module_finetune import (
+    DataPreprocessor,
+    load_backbone,
+    seed_all,
+)
+
+from interpretability.utils import (
+    MatrixRowDataset,
+    NpzCellDataset,
+    build_module_metadata,
+    ensure_dir,
+    load_module_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("--performer_checkpoint", type=str, required=True)
+    p.add_argument("--config_pkl", type=str, required=True)
+    p.add_argument("--module_data_dir", type=str, default="../module_data_all")
+    p.add_argument("--module_json", type=str, default=None)
+    p.add_argument("--gene_list_path", type=str, default="../../selected_genes_27k.txt")
+    p.add_argument("--gene_mapping_path", type=str, default="../../gene2ens.csv")
+    p.add_argument("--gene2vec_data_path", type=str, default="../../selected_gene2vec_27k.npy")
+    p.add_argument("--gene2ens_file", type=str, default="../../gene2ens.csv")
+    p.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
+    p.add_argument("--shuffle_module_genes", action="store_true")
+    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 module_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_module_position_embedding(x_modules)
+        else:
+            x = backbone._add_module_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, module_data, model_config = load_backbone(args, device)
+    backbone.eval()
+
+    num_modules = len(module_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_modules.json")
+    module_names = load_module_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_module_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/module_gene_attention.py

@@ -0,0 +1,1090 @@
+"""ModuleEmbedding Attention-based Gene Interpretability Analysis (Streaming).
+
+Interpretability analysis using single-head content-based attention pooling in ModuleEmbedding.
+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:
+  python module_gene_attention.py \
+      --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
+      --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+      --module_data_dir ../../module_data_all \
+      --gene_list_path ../../../selected_genes_27k.txt \
+      --gene_mapping_path ../../../gene2ens.csv \
+      --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+      --gene2ens_file ../../../gene2ens.csv \
+      --data_path ../../../zero-shot-batch/pancreas_scib.h5ad \
+      --celltype_column celltype \
+      --batch_size 64 \
+      --output_dir ../results/gene_attention/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 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()
+PERFORMER_MODULE = HERE.parents[2]
+REPO_ROOT = HERE.parents[3]
+for p in (str(PERFORMER_MODULE), str(REPO_ROOT)):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+# Publication utilities
+MODULE_IMPORTANCE_DIR = PERFORMER_MODULE / "interpretability" / "module_importance"
+if str(MODULE_IMPORTANCE_DIR) not in sys.path:
+    sys.path.insert(0, str(MODULE_IMPORTANCE_DIR))
+
+from pub_utils import setup_publication_style, savefig as pub_savefig
+
+from cell_annotation.performer_module_finetune import (
+    DataPreprocessor,
+    load_backbone,
+    seed_all,
+)
+from interpretability.utils import (
+    MatrixRowDataset,
+    build_module_metadata,
+    ensure_dir,
+    load_module_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("--performer_checkpoint", type=str, required=True)
+    p.add_argument("--config_pkl", type=str, required=True)
+    p.add_argument("--module_data_dir", type=str, default="../../module_data_all")
+    p.add_argument("--module_json", type=str, default=None)
+    p.add_argument("--gene_list_path", type=str, default="../../../selected_genes_27k.txt")
+    p.add_argument("--gene_mapping_path", type=str, default="../../../gene2ens.csv")
+    p.add_argument("--gene2vec_data_path", type=str, default="../../../selected_gene2vec_27k.npy")
+    p.add_argument("--gene2ens_file", type=str, default="../../../gene2ens.csv")
+    p.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
+    p.add_argument("--shuffle_module_genes", action="store_true")
+    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, module_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/module_geometry.py

@@ -0,0 +1,654 @@
+"""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_modules.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 -m performer_module.interpretability.module_geometry \\
+    --performer_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()
+PERFORMER_MODULE = HERE.parents[1]
+REPO_ROOT = HERE.parents[2]
+for p in (str(PERFORMER_MODULE), str(REPO_ROOT)):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+from cell_annotation.performer_module_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_module_metadata,
+    ensure_dir,
+    load_module_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.performer_module_finetune)
+    parser.add_argument("--performer_checkpoint", type=str, required=True)
+    parser.add_argument("--config_pkl", type=str, required=True)
+    parser.add_argument("--module_data_dir", type=str, default="../module_data_all")
+    parser.add_argument("--module_json", type=str, default=None,
+                        help="Path to merged_modules.json. Defaults to <module_data_dir>/merged_modules.json.")
+    parser.add_argument("--gene_list_path", type=str, default="../../selected_genes_27k.txt")
+    parser.add_argument("--gene_mapping_path", type=str, default="../../gene2ens.csv")
+    parser.add_argument("--gene2vec_data_path", type=str, default="../../selected_gene2vec_27k.npy")
+    parser.add_argument("--gene2ens_file", type=str, default="../../gene2ens.csv")
+    parser.add_argument("--gene2ens_column", type=str, default="Ensembl gene ID")
+    parser.add_argument("--shuffle_module_genes", action="store_true",
+                        help="Randomly reshuffle gene-to-module assignment (control / ablation).")
+    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_module_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_modules.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_module_json(args: argparse.Namespace) -> str:
+    """Resolve relative paths against performer_module/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_modules.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, module_data, model_config = load_backbone(args, device)
+    backbone.eval()
+    base_dim = int(model_config["base_dim"])
+    num_modules = len(module_data)
+    logger.info("Loaded backbone: num_modules=%d, base_dim=%d", num_modules, base_dim)
+
+    # 2. Load module names (must align with module_data order)
+    module_json_path = resolve_module_json(args)
+    module_names = load_module_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_module_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_module_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/module_importance/README.md

@@ -0,0 +1,311 @@
+# 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
+
+```bash
+cd ../data/250010166/scLong/performer_module/interpretability/module_importance
+torchrun --nproc_per_node=4
+python module_importance.py \
+    --finetune_checkpoint mae_tost_48M_cross_attn_only_pure_memory/attn_pool_full_n1_lrh0.0001_lrb0.0001_seed2021/best_model.pt \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/Zheng68K.h5ad \
+    --label_column celltype \
+    --mask_mode post_encoder \
+    --batch_size 128 \
+    --top_k 10 20 \
+    --output_dir ../results/importance/zheng68k_post
+```
+
+Pre-encoder mode, more accurate but slower:
+
+```bash
+python module_importance.py \
+    --finetune_checkpoint mae_tost_48M_cross_attn_only_pure_memory/attn_pool_full_n1_lrh0.0001_lrb0.0001_seed2021/best_model.pt \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/Zheng68K.h5ad \
+    --label_column celltype \
+    --mask_mode pre_encoder \
+    --importance_cells 500 \
+    --batch_size 256 \
+    --top_k 10 20 \
+    --output_dir ../results/importance/zheng68k_pre
+```
+
+Pre-encoder mode on all cells with DDP multi-GPU acceleration:
+
+```bash
+torchrun --nproc_per_node=4 module_importance.py \
+    --finetune_checkpoint mae_tost_48M_cross_attn_only_pure_memory/attn_pool_full_n1_lrh0.0001_lrb0.0001_seed2021/best_model.pt \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/Zheng68K.h5ad \
+    --label_column celltype \
+    --mask_mode pre_encoder \
+    --batch_size 512 \
+    --top_k 10 20 \
+    --output_dir ../results/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 masking_ablation.py \
+    --finetune_checkpoint mae_tost_48M_cross_attn_only_pure_memory/attn_pool_full_n1_lrh0.0001_lrb0.0001_seed2021/best_model.pt \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/Zheng68K.h5ad \
+    --label_column celltype \
+    --importance_matrix ../results/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 ../results/importance/ablation/
+```
+
+### 3-B: Permutation Tests
+
+```bash
+# Mode A: shuffle modules, GPU and slow.
+python random_baseline.py \
+    --mode shuffle_modules \
+    --finetune_checkpoint mae_tost_48M_cross_attn_only_pure_memory/attn_pool_full_n1_lrh0.0001_lrb0.0001_seed2021/best_model.pt \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/Zheng68K.h5ad \
+    --label_column celltype \
+    --importance_matrix ../results/importance/zheng68k_pre_13690/importance_matrix.npy\
+    --n_permutations 10 \
+    --output_dir ../results/importance/permutation/
+
+# Mode B: shuffle labels, CPU and fast, requires per_cell_deltas from ablation.
+python random_baseline.py \
+    --mode shuffle_labels \
+    --importance_matrix ../results/importance/zheng68k_pre_13690/importance_matrix.npy \
+    --per_cell_deltas ../results/importance/ablation/per_cell_deltas.npy \
+    --per_cell_labels ../results/importance/ablation/per_cell_labels.npy \
+    --n_permutations 1000 \
+    --output_dir ../results/importance/permutation/
+```
+
+### 3-C: Marker Gene Verification, CPU Only
+
+```bash
+python marker_gene_verification.py \
+    --importance_matrix ../results/importance/zheng68k_pre_13690/importance_matrix.npy \
+    --importance_csv <importance_matrix.csv> \
+    --module_data_dir ../../module_data_all \
+    --output_dir ../results/importance/marker_verification/
+```
+
+### Publication Figures
+
+```bash
+python module_importance_publication.py \
+    --results_dir ../results/importance/zheng68k_pre_13690/ \
+    --ablation_dir ../results/importance/ablation \
+    --permutation_dir ../results/importance/permutation/ \
+    --module_data_dir ../../module_data_all \
+    --output_dir ../results/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 family_importance_analysis.py \
+    --results_dir ../results/importance/zheng68k_pre \
+    --output_dir ../results/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 family_ablation.py \
+    --finetune_checkpoint mae_tost_48M_cross_attn_only_pure_memory/attn_pool_full_n1_lrh0.0001_lrb0.0001_seed2021/best_model.pt \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/Zheng68K.h5ad \
+    --label_column celltype \
+    --batch_size 256 \
+    --output_dir ../results/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 cross_class_similarity.py \
+    --results_dir ../results/importance/zheng68k_pre \
+    --output_dir ../results/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 performer_module_finetune.py \
+    --config_pkl ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory.pkl \
+    --performer_checkpoint ../../mae_tost_48M_cross_attn_only_pure_memory/mae_tost_48M_cross_attn_only_pure_memory_epoch6.pth \
+    --module_data_dir ../../module_data_all \
+    --gene_list_path ../../../selected_genes_27k.txt \
+    --gene_mapping_path ../../../gene2ens.csv \
+    --gene2vec_data_path ../../../selected_gene2vec_27k.npy \
+    --gene2ens_file ../../../gene2ens.csv \
+    --data_path ../../../cell_annotation/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 ../results/family_exp
+```
+
+The same command can be extended with:
+
+- `--use_family_emb` to enable family-type embedding, adding `4 x 200 = 800` parameters.
+- `--use_module_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 |
+| --- | --- |
+| `module_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 |
+| `module_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/module_importance/RESULTS.md

@@ -0,0 +1,221 @@
+# 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: PerformerModuleCellClassifier with attention-pooling head and 48M backbone.
+- Baseline test accuracy: **83.1%**.
+
+## Data Sources
+
+All analyses are based on the following core files, generated by `module_importance.py`
+with `pre_encoder` over all 13,690 test cells:
+
+| File | Path | Description |
+| --- | --- | --- |
+| `importance_matrix.npy` | `results/importance/zheng68k_pre_13690/importance_matrix.npy` | `(11 x 660)` DeltaLogit matrix |
+| `importance_matrix.csv` | `results/importance/zheng68k_pre_13690/importance_matrix.csv` | Same matrix with row and column labels |
+| `enrichment_results.csv` | `results/importance/zheng68k_pre_13690/enrichment_results.csv` | Enrichment-test results |
+| `top_modules_per_celltype.json` | `results/importance/zheng68k_pre_13690/top_modules_per_celltype.json` | Top-module details |
+| `merged_modules.json` | `module_data_all/merged_modules.json` | Gene composition of 660 modules |
+
+Marker verification uses `importance_matrix.npy`, `importance_matrix.csv`, and
+`merged_modules.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
+
+![Importance Heatmap (clustered)](publication_figures/fig4_importance_heatmap_clustered.png)
+
+---
+
+## 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.
+
+![Ablation Curves](publication_figures/fig1_ablation_curves.png)
+
+![Per-class Ablation](publication_figures/fig_s1_ablation_per_class.png)
+
+---
+
+## 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.
+
+![Permutation Null](publication_figures/fig2_permutation_null.png)
+
+---
+
+## 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.
+
+![Marker Ranks](marker_verification/fig_marker_ranks.png)
+
+![Marker Summary Bar](marker_verification/fig_marker_summary_bar.png)
+
+---
+
+## 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%** |
+
+![Enrichment Summary](publication_figures/fig5_enrichment_summary.png)
+
+---
+
+## 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/module_importance/TODO.md

@@ -0,0 +1,68 @@
+# 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] `module_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.