Remove reverse perturbation code and refresh README documentation

Remove the ten reverse perturbation implementation and protocol files. Update the project README and source README, including their repository trees. Model weights, paired configuration, datasets, and unrelated source files are unchanged.

BioMod/README.md

@@ -44,20 +44,6 @@ BioMod/                                                   # Task code: training,
 ├── 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

BioMod/gene_inverse_norman/README.md

@@ -1,217 +0,0 @@
-# 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

@@ -1,300 +0,0 @@
-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

@@ -1,660 +0,0 @@
-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

@@ -1,405 +0,0 @@
-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

@@ -1,837 +0,0 @@
-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

@@ -1,20 +0,0 @@
-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

@@ -1,59 +0,0 @@
-{
-  "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

@@ -1,20 +0,0 @@
-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

@@ -1,61 +0,0 @@
-{
-  "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

@@ -1,166 +0,0 @@
-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

README.md

@@ -10,38 +10,165 @@ library_name: pytorch
 
 # BioMod
 
-Single-cell transcriptomics foundation model (BioMod / PerformerModule architecture),
-released by **CTC-Atlas** — the joint laboratory of **Raymind Biotechnology Co., Ltd** 
-and the **Shenzhen Research Institute of Big Data (SRIBD)**.
+BioMod is a single-cell transcriptomics foundation model built around gene modules
+and a Performer-based encoder. This repository provides pretrained model assets,
+pretraining and fine-tuning code, downstream task tutorials, and interpretability
+tools.
 
-The current local checkpoint is `weight/0810model/`. Start with the
-[0810 quick start](#0810-quick-start) for a complete example using its paired
-weights and configuration. The outer `BioMod/` directory is the project root;
-the inner `BioMod/` directory contains the source code.
+BioMod is released by **CTC-Atlas**, the joint laboratory of **Raymind Biotechnology
+Co., Ltd** and the **Shenzhen Research Institute of Big Data (SRIBD)**.
 
-## Task navigation
+[Installation](#installation) · [Pretrained model](#pretrained-model) ·
+[Tasks and examples](#tasks-and-examples) · [Usage notes](#usage-notes) ·
+[Repository layout](#repository-layout)
 
-| Task | Entry point | Guide or tutorial |
+## Installation
+
+Download the project from [Hugging Face](https://huggingface.co/ctcatlas/BioMod).
+Run the following commands from the parent of your intended project directory:
+
+```bash
+python -m pip install --upgrade huggingface_hub
+hf auth login
+hf download ctcatlas/BioMod --local-dir ./BioMod
+cd BioMod
+```
+
+The download includes source code, model dependencies, the pretrained checkpoint,
+and datasets. To omit datasets, add `--exclude "data/**"` to the download command.
+If you already have the project, start from its root directory.
+
+Create and activate the supplied Conda environment:
+
+```bash
+conda env create -f BioMod/environment.yml
+conda activate performer_module
+```
+
+The environment specification uses Python 3.9, PyTorch 2.6, CUDA 12.4, and PyTorch
+Geometric 2.6. See [environment.yml](BioMod/environment.yml) and
+[requirements.txt](BioMod/requirements.txt) for the complete dependency lists.
+Match the GPU packages to your system when adapting the environment.
+
+Pretraining and new-data reconstruction evaluation use CUDA/NCCL. The zero-shot
+runner selects CUDA when available and otherwise uses CPU. Native Windows
+execution has not been validated for the supplied environment.
+
+## Pretrained model
+
+The **0810 checkpoint** is included in this release. Use its weights and
+configuration together; the configuration records the model architecture and
+gene-module definitions.
+
+| Asset | File |
+| --- | --- |
+| Model weights | [mae_tost_70M_75mask_ec2_epoch0.pth](weight/0810model/mae_tost_70M_75mask_ec2_epoch0.pth) |
+| Model configuration | [mae_tost_70M_75mask_ec2.pkl](weight/0810model/mae_tost_70M_75mask_ec2.pkl) |
+| Model gene order | [selected_genes_27k.txt](model_dependencies/selected_genes_27k.txt) |
+| Gene identifier mapping | [gene2ens.csv](model_dependencies/gene2ens.csv) |
+| Gene2Vec embeddings | [selected_gene2vec_27k.npy](model_dependencies/selected_gene2vec_27k.npy) |
+| Shared module definitions | [merged_modules.json](model_dependencies/module_data_all/merged_modules.json) |
+
+The model components are bundled under
+[`model_dependencies/performer_pytorch_cont/`](model_dependencies/performer_pytorch_cont/).
+Task-specific checkpoint arguments are listed under
+[Usage notes](#checkpoint-arguments).
+
+## Tasks and examples
+
+The following tasks have implementations in this repository. The links identify
+available source code and examples; these workflows have not been validated end
+to end with the bundled checkpoint.
+
+| Task | Implementation | Notebook or guide |
 | --- | --- | --- |
 | Cell-type annotation | [Fine-tuning script](BioMod/cell_annotation/performer_module_finetune.py) | [Notebook](BioMod/toturials/cell_annotation.ipynb) |
-| Zero-shot batch integration | [Zero-shot runner](BioMod/zero_shot_batch/performer_module_zero_shot.py) | [0810 quick start](#0810-quick-start) |
-| Single-drug response | [Training script](BioMod/drug_response/train_drug_performer_module.py) | [Notebook](BioMod/toturials/drug_response_prediction.ipynb) |
-| Drug-combination synergy | [Training script](BioMod/drug_combination/train_drug_combination_performer_module.py) | [Task guide](BioMod/drug_combination/README.md), [notebook](BioMod/toturials/drug_combination.ipynb) |
-| Gene perturbation response | [Training script](BioMod/gene_perturbation/train_module.py) | [Task guide](BioMod/gene_perturbation/README.md), [notebook](BioMod/toturials/gene_perturbation.ipynb) |
-| Reverse perturbation prediction | [Benchmark builder](BioMod/gene_inverse_norman/build_norman_gene_inverse_benchmark.py), [training](BioMod/gene_inverse_norman/finetune_gene_inverse.py), [evaluation](BioMod/gene_inverse_norman/evaluate_gene_inverse.py) | [Task guide](BioMod/gene_inverse_norman/README.md) |
-| Reconstruction evaluation on new data | [Evaluation script](BioMod/eval_new_data/evaluate_new_data.py) | [Launcher examples](BioMod/eval_new_data/cmd.txtx) |
-| Interpretability | [Module geometry](BioMod/interpretability/module_geometry.py) | [Overview](BioMod/interpretability/README.md), [attention](BioMod/interpretability/attention/README.md), [module importance](BioMod/interpretability/module_importance/README.md), [text alignment](BioMod/interpretability/module_text_alignment/DESIGN.md) |
-| Pretraining and continued pretraining | [Pretraining script](BioMod/performer_module.py) | [Launcher example](BioMod/command.txt), [source guide](BioMod/README.md) |
-
-Task guides and notebooks include historical checkpoint and server paths.
-Use the current asset paths below and read [Notebook setup](#notebook-setup)
-before following a tutorial. The launcher examples are references to adapt to
-your environment, not additional steps in the quick start.
+| Zero-shot cell embeddings and batch-effect evaluation | [Embedding and evaluation script](BioMod/zero_shot_batch/performer_module_zero_shot.py) | [Notebook](BioMod/toturials/zero_shot_batch_integration.ipynb) |
+| Single-drug response regression (IC50) | [Training script](BioMod/drug_response/train_drug_performer_module.py) | [Notebook](BioMod/toturials/drug_response_prediction.ipynb) |
+| Drug-pair synergy classification | [Training script](BioMod/drug_combination/train_drug_combination_performer_module.py) | [Task guide](BioMod/drug_combination/README.md), [notebook](BioMod/toturials/drug_combination.ipynb) |
+| Gene perturbation response prediction | [Training script](BioMod/gene_perturbation/train_module.py) | [Task guide](BioMod/gene_perturbation/README.md), [notebook](BioMod/toturials/gene_perturbation.ipynb) |
+
+### Pretraining, evaluation, and analysis scripts
+
+The repository also includes the following supporting workflows.
+
+| Workflow | Implementation | Documentation |
+| --- | --- | --- |
+| Pretraining and continued pretraining | [Pretraining script](BioMod/performer_module.py) | [Source guide](BioMod/README.md), [launcher reference](BioMod/command.txt) |
+| Expression reconstruction evaluation on new data | [Evaluation script](BioMod/eval_new_data/evaluate_new_data.py) | [Launcher reference](BioMod/eval_new_data/cmd.txtx) |
+| Module interpretability analysis | [Module geometry](BioMod/interpretability/module_geometry.py), [module importance](BioMod/interpretability/module_importance/module_importance.py), [text alignment](BioMod/interpretability/module_text_alignment/module_text_alignment.py) | [Overview](BioMod/interpretability/README.md), [attention](BioMod/interpretability/attention/README.md), [module importance guide](BioMod/interpretability/module_importance/README.md) |
+
+Tutorials and launcher references retain some historical checkpoint paths,
+server paths, and arguments. Adapt them to the current assets above and check
+the selected script's `--help` before running. Some notebook plotting steps
+also refer to external scripts or reference results that are not bundled here.
+
+## Usage notes
+
+### Project paths and notebooks
+
+The project root contains `BioMod/`, `model_dependencies/`, `weight/`, and
+`data/`. The inner `BioMod/` directory contains the source code. Task commands
+may assume different working directories; adjust relative paths when following
+a guide.
+
+Before running a notebook, update its configuration cells:
+
+| Setting | Location |
+| --- | --- |
+| `BASE` or `ASSET_ROOT` | Project root |
+| `PROJECT_ROOT` | Inner `BioMod/` source directory |
+| `MODULE_MODEL` or `LEGACY_MODULE_MODEL` | `model_dependencies/` under the project root |
+| Weight and configuration paths | The paired 0810 files listed above |
+| `MODULE_DATA_DIR` | `BioMod/module_data_all/` under the project root |
+| Data and output paths | Your task inputs and a writable output directory |
+
+Checkpoint variable names differ by notebook: `CKPT_PATH` / `CONFIG_PKL`
+for cell annotation, drug response, and zero-shot integration;
+`SCFM_CKPT` / `SCFM_HYPER` for gene perturbation; and
+`CHECKPOINT_PATH` / `HYPER_PARAMS_PATH` for drug combination.
+Updating environment variables alone does not replace hard-coded notebook paths.
+
+### Checkpoint arguments
+
+| Task | Weights argument | Configuration argument |
+| --- | --- | --- |
+| Drug response / drug combination | `--checkpoint_path` | `--hyper_params_path` |
+| Zero-shot / gene perturbation | `--scfm_ckpt_path` | `--scfm_hyper_params_path` |
+| Cell annotation / backbone interpretation | `--performer_checkpoint` | `--config_pkl` |
+| Continued pretraining | `--model_path` | `--model_config_path` or paired-file discovery |
+
+New-data reconstruction evaluation has a separate checkpoint-directory interface;
+refer to its script for the corresponding arguments.
+
+### Data and evaluation
+
+Datasets are stored under `data/`. Input formats, metadata fields, preprocessing,
+and train/validation splits depend on the task; follow the linked implementation
+and tutorial.
+
+For zero-shot embeddings, provide an h5ad expression matrix and specify the gene,
+batch, and cell-type metadata fields. The runner aligns genes to the model gene
+list and fills missing genes with zeros. It does not normalize or log-transform
+the input matrix, so prepare expression values according to the intended protocol.
+Resolve duplicate gene identifiers and inspect gene overlap before evaluating
+embeddings.
+
+The drug-response and drug-combination workflows use validation metrics for
+checkpoint selection. Their validation results should not be presented as
+performance on an independent final test set.
+
+For checks of the shared training utilities, see
+[Training utility checks](BioMod/README.md#training-utility-checks).
 
 ## Repository layout
 
-Current local checkout, excluding the top-level `data/` and `.cache/` directories.
-Module definitions, model dependencies, and the remaining checkpoint are included.
+The tree below includes source files, model dependencies, and the pretrained
+checkpoint. The top-level `data/` and `.cache/` directories are omitted for
+readability. Folder descriptions explain where to find each component.
+
+<details>
+<summary>Expand the full repository structure</summary>
 
 ```text
 BioMod/                                                   # Project root: source code, model dependencies and weights
@@ -62,20 +189,6 @@ BioMod/                                                   # Project root: source
 │   ├── 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
@@ -181,375 +294,9 @@ BioMod/                                                   # Project root: source
 └── README.md
 ```
 
-> The local `data/` directory is present but omitted from this view. Only
-> `weight/0810model/` is currently included in the local checkpoint directory.
-
-## Environment requirements
-
-The supplied files describe the original CUDA server environment. They are
-configuration specifications, not a tested compatibility matrix for this local
-checkout.
-
-| Component | Declared configuration |
-| --- | --- |
-| Conda environment | `performer_module` |
-| Python | `3.9` |
-| PyTorch | `2.6.*` in Conda; `2.6.0+cu124` in the pip requirements |
-| CUDA runtime | `pytorch-cuda=12.4`; CUDA 12.4 wheel sources in the pip requirements |
-| TorchVision / TorchAudio | `0.21.*` / `2.6.*` |
-| PyTorch Geometric | `2.6.*` |
-| Notebook tools | JupyterLab and ipykernel |
-
-See [environment.yml](BioMod/environment.yml) and
-[requirements.txt](BioMod/requirements.txt) for the full dependency lists.
-The pip requirements also pin GPU packages; review those pins before applying
-them to an existing environment.
-
-Pretraining uses the local [training utilities](BioMod/training_utils.py) for
-seeding, distributed sampling and reduction, learning-rate scheduling, and
-checkpoint saving. See the [utility checks](BioMod/README.md#training-utility-checks)
-for CPU tests and an entry-point import check.
-
-The zero-shot runner selects CUDA when available and otherwise selects CPU.
-Other tasks have different requirements: pretraining and new-data reconstruction
-evaluation use CUDA/NCCL. Native Windows support, CPU performance, GPU memory
-requirements, and elapsed time have not been validated here.
-
-### Working directory and installation
-
-All commands in the quick start run from the **outer project root**, which
-contains `BioMod/`, `model_dependencies/`, `weight/`, and `data/`.
-Choose the working-directory command for your shell:
-
-```powershell
-# PowerShell on this Windows checkout
-Set-Location "E:/SLAI/BioMod"
-```
-
-```bash
-# Bash on Linux or WSL; adjust the mount path if necessary
-cd /mnt/e/SLAI/BioMod
-```
-
-With Conda available in that shell, create and activate the environment:
-
-```bash
-conda env create -f BioMod/environment.yml
-conda activate performer_module
-```
-
-These two Conda commands use the same syntax in Bash and PowerShell. Keep the
-working directory at the outer root. Environment installation and model inference
-have not been executed to validate these examples.
-
-## 0810 quick start
-
-Use the following two files together:
-
-| Asset | Path relative to the project root |
-| --- | --- |
-| Weights | `weight/0810model/mae_tost_70M_75mask_ec2_epoch0.pth` |
-| Configuration | `weight/0810model/mae_tost_70M_75mask_ec2.pkl` |
-
-The 0810 release is documented as a 2-layer encoder with a 27,874-gene axis
-and 660 modules. The paired configuration defines the exact architecture and
-module layout.
-The loader restores those settings and uses the shared gene assets under
-`model_dependencies/`. Do not substitute the configuration from the old 0629
-or 0407 examples. The `epoch0` filename alone does not establish performance
-on any downstream task.
-
-### Small zero-shot evaluation
-
-After activating the environment, choose **one** command below. Both commands
-use the same local lung dataset, 512 cells, mean pooling, and a batch size of 1.
-The dataset metadata contains 21,955 cells and the required `gene_name`, `sample`,
-and `cell_type` fields. This metadata check does not establish that the expression
-preprocessing or checkpoint performance is suitable for a new benchmark.
-
-**Bash:**
-
-```bash
-python BioMod/zero_shot_batch/performer_module_zero_shot.py \
-  --scfm_ckpt_path weight/0810model/mae_tost_70M_75mask_ec2_epoch0.pth \
-  --scfm_hyper_params_path weight/0810model/mae_tost_70M_75mask_ec2.pkl \
-  --gene2vec_path model_dependencies/selected_gene2vec_27k.npy \
-  --scfm_genes_file model_dependencies/selected_genes_27k.txt \
-  --gene_mapping_path model_dependencies/gene2ens.csv \
-  --module_data_path model_dependencies/module_data_all/merged_modules.json \
-  --data_path data/downstream/final_data/Kim2020_Lung_noMalignant.h5ad \
-  --dataset_preset lung \
-  --gene_col gene_name \
-  --batch_key sample \
-  --label_key cell_type \
-  --n_cells 512 \
-  --batch_size 1 \
-  --metrics_mode simple \
-  --pool_method mean \
-  --output_dir outputs/0810_lung_smoke
-```
-
-**PowerShell:**
-
-```powershell
-python BioMod/zero_shot_batch/performer_module_zero_shot.py `
-  --scfm_ckpt_path weight/0810model/mae_tost_70M_75mask_ec2_epoch0.pth `
-  --scfm_hyper_params_path weight/0810model/mae_tost_70M_75mask_ec2.pkl `
-  --gene2vec_path model_dependencies/selected_gene2vec_27k.npy `
-  --scfm_genes_file model_dependencies/selected_genes_27k.txt `
-  --gene_mapping_path model_dependencies/gene2ens.csv `
-  --module_data_path model_dependencies/module_data_all/merged_modules.json `
-  --data_path data/downstream/final_data/Kim2020_Lung_noMalignant.h5ad `
-  --dataset_preset lung `
-  --gene_col gene_name `
-  --batch_key sample `
-  --label_key cell_type `
-  --n_cells 512 `
-  --batch_size 1 `
-  --metrics_mode simple `
-  --pool_method mean `
-  --output_dir outputs/0810_lung_smoke
-```
-
-The flags and input paths were checked against this checkout; checkpoint loading
-and inference remain untested. The PowerShell block documents shell syntax, not
-validated native-Windows dependency compatibility.
-
-`--module_data_path` is a fallback when the paired configuration lacks module
-sets; it does not override module sets already stored in that configuration.
-For a different dataset, replace the input path and metadata fields as described
-under [Input and output](#input-and-output).
-
-`--n_cells 512` subsamples after loading the entire h5ad file. For another input,
-choose a positive value no larger than the number of cells remaining after
-filtering, or use `--n_cells all`. The limit does not reduce the initial file-loading
-memory requirement. A small subset can lack the batch or cell-type diversity
-needed by the metrics.
-
-For a full evaluation, use `--n_cells all` and `--metrics_mode scgpt`, and select
-a new output directory. The `simple` mode computes only the available batch
-silhouette metric; its subset results are not a full benchmark.
-
-### Check the result
-
-A completed run exits successfully and prints `Zero-shot evaluation completed!`.
-The two main files are listed under [Expected outputs](#expected-outputs).
-After running the 512-cell example, this optional command checks their basic
-structure without loading model weights:
-
-```bash
-python -c "from pathlib import Path; import numpy as np; import pandas as pd; p=Path('outputs/0810_lung_smoke'); x=np.load(p/'sclong_module_batch_cell_emb_mode_ce.npy', allow_pickle=False); m=pd.read_csv(p/'sclong_module_batch_cell_emb_mode.csv'); assert x.ndim==2 and x.shape[0]==512 and x.shape[1]>0; assert np.isfinite(x).all(); assert {'metric','label','value'}.issubset(m.columns); print('Embedding shape:', x.shape, 'Metric rows:', len(m))"
-```
-
-The same one-line check works in PowerShell. Adjust the expected cell count if
-you change the command. Passing these checks establishes output structure, not
-biological accuracy or a performance claim. A metric CSV with no data rows can
-occur when the subset contains only one batch or no valid metric values.
-
-### Checkpoint arguments for other tasks
-
-| Task | Weights argument | Configuration argument |
-| --- | --- | --- |
-| Drug response / drug combination | `--checkpoint_path` | `--hyper_params_path` |
-| Zero-shot / gene perturbation / gene inverse | `--scfm_ckpt_path` | `--scfm_hyper_params_path` |
-| Cell annotation / backbone interpretation | `--performer_checkpoint` | `--config_pkl` |
-| Continued pretraining | `--model_path` | `--model_config_path` or automatic paired-file discovery |
-
-Use the 0810 pair above when adapting these task commands. New-data reconstruction
-evaluation uses its own checkpoint-directory interface; inspect its linked
-launcher and parser rather than applying this table to that entry point.
-
-## Notebook setup
-
-The notebooks in [`BioMod/toturials/`](BioMod/toturials/) retain historical
-`0407weight`, `module_model`, and `/data/...` settings. Review the first
-configuration cells before executing them. Updating an environment variable alone
-does not update every notebook: the drug-combination notebook reads
-`PERFORMER_MODULE_ROOT` and `BIOMODULEFM_ASSET_ROOT`, while the others contain
-hard-coded paths.
-
-| Setting | Local meaning |
-| --- | --- |
-| `BASE` or `ASSET_ROOT` | Outer project root, for example `E:/SLAI/BioMod` in native Windows Python |
-| `PROJECT_ROOT` | Inner source directory, for example `E:/SLAI/BioMod/BioMod` |
-| `MODULE_MODEL` or `LEGACY_MODULE_MODEL` | Shared assets directory, `model_dependencies/` under the outer root |
-| Weight and configuration variables | The exact 0810 `.pth` and `.pkl` filenames listed above |
-| `MODULE_DATA_DIR` | `BioMod/module_data_all/` under the outer root |
-| Data and output variables | Task-specific inputs and a writable output directory |
-| Python executable and launch commands | The active environment and the syntax supported by your shell |
-
-The variable names for the paired model files differ by notebook:
-
-- Cell annotation, drug response, and zero-shot: `CKPT_PATH` and `CONFIG_PKL`.
-- Gene perturbation: `SCFM_CKPT` and `SCFM_HYPER`.
-- Drug combination: `CHECKPOINT_PATH` and `HYPER_PARAMS_PATH`.
-
-Also replace derived paths such as `WEIGHT_DIR`, `DATA_PATH` / `DATA_DIR` /
-`DATA_ROOT`, and the notebook's output root. Confirm the path-check cells before
-turning on `RUN_TRAINING` where that switch is provided. Under WSL, use Linux
-mount paths such as `/mnt/e/SLAI/BioMod` rather than Windows drive paths.
-
-Some plotting cells require additional adaptation. The zero-shot notebook calls
-`module_model/draw/zero_shot_umap_scib.py`, which is not included in this checkout;
-use the bundled CLI above for the documented 0810 run. The drug-response notebook
-also refers to external historical plotting scripts and reference results.
-Changing only the checkpoint paths does not supply those external files.
-
-## Input and output
-
-### Zero-shot input requirements
-
-The quick start uses
-`data/downstream/final_data/Kim2020_Lung_noMalignant.h5ad`.
-The data directory is intentionally omitted from the repository tree above.
-Other tasks have different input formats; follow their linked task guides.
-
-| Input | Requirement or behavior |
-| --- | --- |
-| Expression matrix | An h5ad file with cells in rows and genes in columns of `adata.X`; dense and sparse matrices are handled |
-| Expression preprocessing | Prepare the matrix according to the checkpoint/task protocol; this runner does not normalize, log-transform, scale, or select highly variable genes |
-| Gene identifiers | Supply the column in `adata.var` with `--gene_col`; with no column selected, the runner uses `adata.var.index` |
-| Gene mapping | Target Ensembl IDs can match directly; gene symbols are mapped using `model_dependencies/gene2ens.csv` |
-| Batch labels | Select an existing `adata.obs` column with `--batch_key` |
-| Cell-type labels | Select one or more existing `adata.obs` columns with `--label_key`; evaluation requires valid labels |
-| Lung preset | Defaults to `gene_name` in `adata.var`, and `sample` / `cell_type` in `adata.obs`; explicit command-line values take precedence |
-| Model gene order | Defined by `model_dependencies/selected_genes_27k.txt`; missing target genes are filled with zeros |
-
-Check the printed gene-resolution summary and gene intersection count before
-interpreting results. Resolve duplicate gene identifiers before inference: this
-runner's index mapping keeps the last matching column rather than aggregating
-duplicates. No minimum acceptable gene-overlap threshold is enforced by the script.
-
-### Expected outputs
-
-The example writes to `outputs/0810_lung_smoke/`:
-
-| File | Contents |
-| --- | --- |
-| `sclong_module_batch_cell_emb_mode.csv` | Evaluation rows with columns `metric`, `label`, and `value` |
-| `sclong_module_batch_cell_emb_mode_ce.npy` | Cell embeddings with shape `(processed_cells, embedding_dimension)` |
-| `sclong_module_raw_module_emb.npy` | Optional module embeddings, written only with `--save_raw_module_emb` |
-
-The legacy output prefixes are the filenames produced by the current script.
-The runner does not export a new h5ad file, cell IDs, or UMAP figures. Embedding
-rows follow the processed AnnData order; after subsampling, do not pair them with
-the full original observation list. If downstream work needs cell IDs, prepare
-and retain a documented input subset and process it with `--n_cells all`.
-
-Metrics are evaluated before the main output files are saved. An evaluation
-error can therefore prevent the final embedding file from being written even
-when embedding extraction finished.
-
-## Troubleshooting
-
-| Symptom | What to check |
-| --- | --- |
-| Missing script, checkpoint, or gene asset | Confirm the outer working directory and the exact paths in the 0810 example; older guides may still refer to 0629 or 0407 assets |
-| Configuration or state-dictionary mismatch | Use the matching `.pth` and `.pkl` pair and the current loader; do not edit serialized weights to resolve names |
-| Missing gene, batch, or label column | Check `adata.var` / `adata.obs` and pass the correct `--gene_col`, `--batch_key`, and `--label_key` |
-| Empty metrics or a small-subset metric error | Check batch and label coverage, then use a larger suitable subset or the full dataset |
-| Invalid cell limit or an empty inference batch | Use `--n_cells all`, or a positive count no greater than the cells available after filtering |
-| Insufficient memory | Keep inference batches small and assess host RAM separately; h5ad loading happens before subsampling |
-| Notebook references an absent `draw/` script | Use the bundled CLI or adapt the external plotting workflow; renaming weight paths alone is insufficient |
-
-## Download
-
-Run these commands from the **parent of the destination directory**. Skip this
-section when using the files already present locally. Authenticate with a
-Hugging Face account that has access to `ctcatlas/BioMod`:
-
-```bash
-hf auth login
-
-# Complete project: source code, model dependencies, 0810 weights and datasets
-hf download ctcatlas/BioMod --repo-type=model --local-dir ./BioMod
-
-# Source code, model dependencies and weights, without the data directory
-hf download ctcatlas/BioMod --repo-type=model --local-dir ./BioMod \
-  --exclude "data/**"
-
-# The 0810 checkpoint and its matching architecture configuration only
-hf download ctcatlas/BioMod \
-  weight/0810model/mae_tost_70M_75mask_ec2_epoch0.pth \
-  weight/0810model/mae_tost_70M_75mask_ec2.pkl \
-  --repo-type=model --local-dir ./BioMod
-```
-
-The complete repository includes the `data/` directory, which is omitted from
-[the repository layout](#repository-layout) for readability. Use the command
-with `--exclude "data/**"` when the benchmark datasets are not needed.
-
-This release uses the `BioMod/` source-directory name shown throughout this
-README. The original [upstream repository](https://huggingface.co/ctcatlas/ModulaFM)
-is retained separately; its historical layout and 0629 checkpoint are not part
-of this release.
-
-## Historical 0629 reference
-
-The 0629 checkpoint and its paired PKL have been removed from this local checkout.
-Only `weight/0810model/` remains. The 0629 examples below are upstream reference
-commands and require restoring that checkpoint/config pair before use.
-
-<details>
-<summary>Original 0629 architecture and Bash examples</summary>
-
-Start these reference blocks from the outer project root. The final pretraining
-block changes into the inner source directory.
-
-The primary release is the 32-layer checkpoint under `weight/0629model/`.
-It uses 27,874 genes, 660 gene modules, a hidden dimension of 1,024, and the
-cross-attention decoder. Use the checkpoint and its PKL as a pair:
-
-```bash
-PRIMARY_MODEL_DIR="$PWD/weight/0629model"
-PRIMARY_MODEL_CKPT="$PRIMARY_MODEL_DIR/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth"
-PRIMARY_MODEL_CONFIG="$PRIMARY_MODEL_DIR/mae_tost_48M_cross_attn_only_pure_memory.pkl"
-```
-
-The PKL records the exact architecture and module definitions used during
-pretraining. Do not combine a checkpoint with a PKL from another model.
-
-Example zero-shot invocation from the repository root:
-
-```bash
-python BioMod/zero_shot_batch/performer_module_zero_shot.py \
-  --data_path /path/to/input.h5ad \
-  --scfm_hyper_params_path "$PRIMARY_MODEL_CONFIG" \
-  --scfm_ckpt_path "$PRIMARY_MODEL_CKPT" \
-  --gene_mapping_path model_dependencies/gene2ens.csv \
-  --scfm_genes_file model_dependencies/selected_genes_27k.txt \
-  --gene2vec_path model_dependencies/selected_gene2vec_27k.npy \
-  --output_dir BioMod/results/zero_shot_0629
-```
-
-For drug-response and drug-combination scripts, pass the same pair through
-`--checkpoint_path` and `--hyper_params_path`. For cell annotation, use
-`--performer_checkpoint` and `--config_pkl`. The detailed task commands are in
-[source README](BioMod/README.md).
-
-To continue pretraining, pass the primary checkpoint through `--model_path`.
-`performer_module.py` automatically finds the accompanying PKL, restores the
-recorded architecture, and uses its module definitions:
-
-```bash
-cd BioMod
-torchrun --standalone --nproc_per_node=1 performer_module.py \
-  --model_path ../weight/0629model/mae_tost_48M_cross_attn_only_pure_memory_epoch12.pth \
-  --gene_list_path ../model_dependencies/selected_genes_27k.txt \
-  --gene_mapping_path ../model_dependencies/gene2ens.csv \
-  --gene2vec_data_path ../model_dependencies/selected_gene2vec_27k.npy \
-  --data_path /path/to/pretraining_data \
-  --data_format h5ad
-```
-
-For a checkpoint whose PKL does not follow the same naming convention, pass
-`--model_config_path /path/to/config.pkl` explicitly. The bundled
-`model_dependencies/performer_pytorch_cont` package is discovered
-automatically; callers no longer need to add it to `PYTHONPATH`.
-
 </details>
 
 ## Partners
 
-- **Raymind Biotechnology Co., Ltd** (睿迈生物) — https://en.raymind.com/
-- **Shenzhen Research Institute of Big Data** (SRIBD) — https://www.sribd.cn/en
+- [Raymind Biotechnology Co., Ltd](https://en.raymind.com/)
+- [Shenzhen Research Institute of Big Data (SRIBD)](https://www.sribd.cn/en)