6.1 Geneformer(AIVC 课程)

# ============================================================
# Geneformer in silico KO pipeline — 全局配置
# 所有可调参数集中于此;改配置即可,不用改代码。
# 路径均相对本文件所在目录(pipeline/)解析,亦支持绝对路径。
# ============================================================

project:
  name: geneformer_norman2019          # 项目名,用于输出前缀与日志命名
  seed: 123                            # 抽样随机种子(固定,保证可复现)
  work_dir: ../workspace               # 运行时工作目录(中间产物全在这)
  log_dir: ../workspace/logs           # 日志目录

paths:
  # ---- 数据:GPU 机器上不存在该文件时,由 init.sh 尝试按 download.data_urls 下载 ----
  data_file: ../data/Norman_2019.h5ad

  # ---- 模型权重:本地优先;不存在或校验失败时,由 init.sh 从 HF 下载 ----
  model_dir: ../models/Geneformer-V2-104M

  # ---- 官方 Geneformer 源码目录(init.sh 自动 clone 到这里) ----
  geneformer_repo: third_party/Geneformer

download:
  # ---- Norman_2019.h5ad 下载源(按顺序逐个尝试,留空则跳过下载,仅提示手动放置) ----
  # 填法示例:
  #   data_urls:
  #     - "https://example.com/Norman_2019.h5ad"
  data_urls: []
  # 期望 md5(可选;填了则下载后校验,不匹配会报错而不是静默使用)
  data_md5: ""
  # 期望最小字节数:小于该值视为下载失败/截断
  data_min_bytes: 100000000

  # ---- 模型权重:从 HuggingFace Hub 拉取 ----
  hf_model_repo: ctheodoris/Geneformer  # 官方模型仓库
  hf_model_subdir: Geneformer-V2-104M   # 仓库内子目录名
  # 完整性校验:真实权重必须大于该字节数(小于此值基本是 git-lfs 指针/损坏文件)
  model_min_bytes: 100000000

data:
  # 快速模式:true=抽样 max_cells 个细胞先跑通全链路;false=全量数据
  fast_mode: true
  max_cells: 5000
  # 可选过滤:只保留指定 cell_type / perturbation_name(留空=不过滤)
  filter_cell_type: null
  filter_perturbation: null
  # obs 列名
  cell_type_key: cell_type
  perturbation_key: perturbation_name

tokenize:
  nproc: 4                             # tokenization 并行进程数
  model_version: V2                    # 词表版本(须与模型一致)
  collapse_gene_ids: true              # 同义 Ensembl ID 合并
  use_h5ad_index: true                 # 用 var.index 作为基因 ID
  custom_attrs:                        # 随 token 保存的元信息列
    cell_type: cell_type
    perturbation_name: perturbation_name

perturb:
  # 扰动基因列表(symbol 或 Ensembl ID 均可);留空=自动从数据+词表交集挑 1 个 TF
  genes_to_perturb: []
  # 自动挑基因时的候选 TF(按顺序取第一个同时存在于数据与词表的)
  candidate_tf_symbols: [GATA1, KLF1, RUNX1, SPI1, CEBPA]
  perturb_type: delete                 # in silico KO:单基因删除
  combos: 0                            # 0=不跑组合扰动
  emb_mode: cls
  emb_layer: 0
  # 以下留空=自动:GPU 上 max_ncells=1000 / forward_batch_size=64 / nproc=4,
  # CPU 上 max_ncells=128 / forward_batch_size=8 / nproc=1
  max_ncells: null
  forward_batch_size: null
  nproc: null

stats:
  mode_name: aggregate_data            # 把多细胞的 cosine shift 聚合
  # 统计哪些基因?留空=继承 perturb.genes_to_perturb
  genes_perturbed: null
  cell_states_to_model: null           # 按 cell_type 分层统计(实验性,留空=全部)

hpc:
  cuda_device: auto                    # auto=自动选第一块 GPU;也可写 "cuda:0"

#!/usr/bin/env bash
# ============================================================
# Geneformer pipeline 一键初始化
# 幂等:可重复执行,已就绪的资源自动跳过。
# 用法: cd pipeline && bash init.sh
# ============================================================
set -euo pipefail

cd "$(dirname "$0")"  # 确保在 pipeline/ 目录下
PIPELINE_DIR="$PWD"

echo "===================================================="
echo " Geneformer pipeline 初始化"
echo "===================================================="

# ---- 1. 拉取官方 Geneformer 源码 ----
GENEFORMER_REPO="third_party/Geneformer"
if [ -d "$GENEFORMER_REPO" ] && [ -f "$GENEFORMER_REPO/geneformer/__init__.py" ]; then
    echo "[init] 源码已就绪: $GENEFORMER_REPO"
else
    echo "[init] 正在 clone 官方 Geneformer 源码 ..."
    mkdir -p third_party
    git clone --depth 1 https://github.com/ctheodoris/Geneformer.git "$GENEFORMER_REPO"
    echo "[init] 源码 clone 完成"
fi

# ---- 2. 检测基础 python 能力(只需要 yaml) ----
python3 -c "import yaml" 2>/dev/null || {
    echo "[init] 需要 PyYAML,尝试安装 ..."
    pip install pyyaml -q 2>/dev/null || pip3 install pyyaml -q
}

# ---- 3. 数据下载 ----
echo "[init] 检查数据 ..."
python3 utils/downloads.py --mode data

# ---- 4. 模型下载 ----
echo "[init] 检查模型权重 ..."
python3 utils/downloads.py --mode model

echo ""
echo "===================================================="
echo " 初始化完成!"
echo " 下一步: bash setup_env.sh   (搭建 conda 环境)"
echo "===================================================="

"""Geneformer in silico KO pipeline — 主入口。

用法(在 pipeline/ 目录下,已激活 geneformer 环境):
    python run_pipeline.py                 # 按顺序跑完所有阶段,已完成的自动跳过
    python run_pipeline.py --stage tokenize   # 从 tokenize 阶段开始(断点续跑)
    python run_pipeline.py --force         # 忽略缓存,全部重跑
    python run_pipeline.py --config /path/to/config.yaml

阶段:
    data     预处理 + 抽样,输出中间 h5ad
    tokenize 生成 Geneformer token 序列(.dataset)
    perturb  in silico KO 扰动(独立子进程,GPU 自适应)
    stats    aggregate_data 统计汇总,输出 CSV
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from utils.config_loader import load_config  # noqa: E402
from utils.logging_utils import get_logger  # noqa: E402

STAGES = ["data", "tokenize", "perturb", "stats"]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Geneformer in silico KO pipeline",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--config", default=None, help="自定义 config.yaml 路径")
    parser.add_argument("--stage", choices=STAGES, default=None,
                        help="从指定阶段开始运行(跳过之前阶段)")
    parser.add_argument("--force", action="store_true",
                        help="忽略已缓存产物,强制重跑所有阶段")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    cfg = load_config(args.config)
    log = get_logger(cfg)
    log.info("配置加载完成: %s", cfg)

    start_idx = STAGES.index(args.stage) if args.stage else 0
    for stage in STAGES[start_idx:]:
        _run_stage(stage, cfg, force=args.force, log=log)

    log.info("全部阶段完成 ✔")
    return 0


def _run_stage(stage: str, cfg: dict, force: bool, log) -> None:
    if stage == "data":
        from steps.prepare_data import run_stage
    elif stage == "tokenize":
        from steps.tokenize import run_stage
    elif stage == "perturb":
        from steps.perturb import run_stage
    elif stage == "stats":
        from steps.summarize import run_stage
    else:  # pragma: no cover
        raise ValueError(f"未知阶段: {stage}")

    log.info("===== 阶段 [%s] 开始 =====", stage)
    run_stage(cfg, force=force, log=log)
    log.info("===== 阶段 [%s] 完成 =====\n", stage)


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env bash
# ============================================================
# Geneformer pipeline 一键环境搭建
# 功能:检测 python/conda → 建 geneformer env → 装 torch CUDA 版
#       → 装 requirements + geneformer 源码 → 自检 CUDA/包版本
# 用法: cd pipeline && bash setup_env.sh
# 可选环境变量:
#   ENV_NAME=geneformer            # conda 环境名
#   TORCH_CUDA_VERSION=cu121       # torch 的 CUDA 版本(默认自动检测)
#   SKIP_TORCH_INSTALL=0           # 1=跳过 torch 安装(已装过时加速)
# ============================================================
set -euo pipefail

cd "$(dirname "$0")"
ENV_NAME="${ENV_NAME:-geneformer}"
PYTHON_VERSION="${PYTHON_VERSION:-3.10}"
SKIP_TORCH_INSTALL="${SKIP_TORCH_INSTALL:-0}"

echo "===================================================="
echo " Geneformer 环境搭建 (env: $ENV_NAME)"
echo "===================================================="

# ---- 1. 定位 conda ----
CONDA_BIN=""
for cand in "$HOME/miniconda3/bin/conda" "$HOME/anaconda3/bin/conda" "$HOME/miniforge3/bin/conda"; do
  if [ -x "$cand" ]; then CONDA_BIN="$cand"; break; fi
done
if [ -z "$CONDA_BIN" ] && command -v conda >/dev/null 2>&1; then
  CONDA_BIN="$(command -v conda)"
fi
if [ -z "$CONDA_BIN" ]; then
  echo "!! 未找到 conda。请先安装 Miniconda3,或把 conda 加入 PATH 后重跑。"
  exit 1
fi
echo "[setup] 使用 conda: $CONDA_BIN"
# shellcheck disable=SC1091
source "$(dirname "$(dirname "$CONDA_BIN")")/etc/profile.d/conda.sh"

# ---- 2. 创建 env(已存在则跳过) ----
if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
  echo "[setup] env '$ENV_NAME' 已存在,直接使用"
else
  echo "[setup] 创建 env '$ENV_NAME' (python=$PYTHON_VERSION) ..."
  conda create -y -n "$ENV_NAME" "python=$PYTHON_VERSION"
fi
conda activate "$ENV_NAME"
echo "[setup] 当前 python: $(which python)"

# ---- 3. 安装 torch (CUDA 版) ----
if [ "$SKIP_TORCH_INSTALL" = "1" ]; then
  echo "[setup] 跳过 torch 安装 (SKIP_TORCH_INSTALL=1)"
elif python -c "import torch; torch.cuda.is_available()" >/dev/null 2>&1; then
  echo "[setup] torch 已安装且可用: $(python -c 'import torch; print(torch.__version__)')"
else
  # 检测 CUDA 版本,给 torch 选择 index-url
  CUDA_VER="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1 || true)"
  echo "[setup] nvidia driver: ${CUDA_VER:-未检测到 nvidia-smi(可能无 GPU 或未加载驱动)}"
  TORCH_CUDA="${TORCH_CUDA_VERSION:-}"
  if [ -z "$TORCH_CUDA" ]; then
    # 默认优先 cu121;新版驱动也可用 cu124。此处给个保守默认,失败会提示手动指定。
    TORCH_CUDA="cu121"
    echo "[setup] 使用默认 torch CUDA 版本: $TORCH_CUDA (可用 TORCH_CUDA_VERSION 覆盖)"
  fi
  pip install torch --index-url "https://download.pytorch.org/whl/${TORCH_CUDA}"
fi

# ---- 4. 安装依赖 + geneformer 源码 ----
echo "[setup] 安装 requirements.txt ..."
pip install -r requirements.txt
if [ -d "third_party/Geneformer" ] && [ -f "third_party/Geneformer/setup.py" ]; then
  echo "[setup] 安装 geneformer 源码 (editable) ..."
  pip install -e third_party/Geneformer
else
  echo "!! 未找到 third_party/Geneformer,跳过源码安装(请先运行 bash init.sh)"
fi

# ---- 5. 自检报告 ----
echo ""
echo "===================================================="
echo " 环境自检报告"
echo "===================================================="
python - <<'PYEOF'
import platform
print(f"python      : {platform.python_version()}")
try:
    import torch
    print(f"torch       : {torch.__version__}")
    print(f"cuda avail  : {torch.cuda.is_available()}")
    print(f"gpu name    : {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}")
except Exception as e:
    print(f"torch       : 导入失败 -> {e}")
for mod in ["transformers", "scanpy", "anndata", "pandas", "datasets"]:
    try:
        m = __import__(mod)
        print(f"{mod:<12}: {getattr(m, '__version__', 'ok')}")
    except Exception as e:
        print(f"{mod:<12}: 导入失败 -> {e}")
try:
    import geneformer
    print("geneformer  : OK")
except Exception as e:
    print(f"geneformer  : 导入失败 -> {e}")
PYEOF
echo ""
echo "===================================================="
echo " 环境就绪!下一步:"
echo "   python run_pipeline.py          # 跑全流程"
echo "   或网页 Notebook 打开 run_pipeline.ipynb"
echo "===================================================="

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any


def load_payload(payload_path: str) -> dict[str, Any]:
    with open(payload_path, 'r', encoding='utf-8') as handle:
        return json.load(handle)


def add_repo_to_path(repo_path: str) -> Path:
    repo = Path(repo_path).resolve()
    if str(repo) not in sys.path:
        sys.path.insert(0, str(repo))
    return repo


def apply_runtime_defaults(payload: dict[str, Any]) -> dict[str, Any]:
    import torch

    use_cuda = torch.cuda.is_available()
    payload = payload.copy()
    if payload.get('max_ncells') is None:
        payload['max_ncells'] = 1000 if use_cuda else 128
    if payload.get('forward_batch_size') is None:
        payload['forward_batch_size'] = 64 if use_cuda else 8
    if payload.get('nproc') is None:
        payload['nproc'] = 4 if use_cuda else 1
    return payload


def run_isp(payload: dict[str, Any]) -> dict[str, Any]:
    add_repo_to_path(payload['geneformer_repo'])
    from geneformer import InSilicoPerturber

    runtime_payload = apply_runtime_defaults(payload)
    Path(runtime_payload['output_directory']).mkdir(parents=True, exist_ok=True)
    isp = InSilicoPerturber(
        perturb_type=runtime_payload.get('perturb_type', 'delete'),
        genes_to_perturb=runtime_payload['genes_to_perturb'],
        combos=runtime_payload.get('combos', 0),
        anchor_gene=runtime_payload.get('anchor_gene'),
        model_type=runtime_payload.get('model_type', 'Pretrained'),
        num_classes=runtime_payload.get('num_classes', 0),
        emb_mode=runtime_payload.get('emb_mode', 'cls'),
        filter_data=runtime_payload.get('filter_data'),
        max_ncells=runtime_payload.get('max_ncells'),
        emb_layer=runtime_payload.get('emb_layer', 0),
        forward_batch_size=runtime_payload.get('forward_batch_size', 8),
        nproc=runtime_payload.get('nproc', 1),
        model_version=runtime_payload.get('model_version', 'V2'),
    )
    isp.perturb_data(
        model_directory=runtime_payload['model_directory'],
        input_data_file=runtime_payload['input_data_file'],
        output_directory=runtime_payload['output_directory'],
        output_prefix=runtime_payload['output_prefix'],
    )
    output_directory = Path(runtime_payload['output_directory'])
    return {
        'status': 'ok',
        'mode': 'isp',
        'output_directory': str(output_directory),
        'output_prefix': runtime_payload['output_prefix'],
        'output_files': sorted([path.name for path in output_directory.glob('*')]),
        'runtime': {
            'max_ncells': runtime_payload['max_ncells'],
            'forward_batch_size': runtime_payload['forward_batch_size'],
            'nproc': runtime_payload['nproc'],
        },
    }



def run_stats(payload: dict[str, Any]) -> dict[str, Any]:
    add_repo_to_path(payload['geneformer_repo'])
    from geneformer import InSilicoPerturberStats

    Path(payload['output_directory']).mkdir(parents=True, exist_ok=True)
    stats = InSilicoPerturberStats(
        mode=payload.get('mode_name', 'aggregate_data'),
        genes_perturbed=payload['genes_perturbed'],
        combos=payload.get('combos', 0),
        anchor_gene=payload.get('anchor_gene'),
        cell_states_to_model=payload.get('cell_states_to_model'),
        model_version=payload.get('model_version', 'V2'),
    )
    stats.get_stats(
        input_data_directory=payload['input_data_directory'],
        null_dist_data_directory=payload.get('null_dist_data_directory'),
        output_directory=payload['output_directory'],
        output_prefix=payload['output_prefix'],
    )
    stats_csv = (Path(payload['output_directory']) / payload['output_prefix']).with_suffix('.csv')
    return {
        'status': 'ok',
        'mode': 'stats',
        'stats_csv': str(stats_csv),
    }



def main() -> int:
    os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')

    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', required=True, choices=['isp', 'stats'])
    parser.add_argument('--payload', required=True)
    args = parser.parse_args()

    payload = load_payload(args.payload)
    if args.mode == 'isp':
        result = run_isp(payload)
    else:
        result = run_stats(payload)

    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())

from setuptools import setup, find_packages

setup(
    name="geneformer",
    version="0.1.0",
    author="Christina Theodoris",
    author_email="<redacted-email>",
    description="Geneformer is a transformer model pretrained \
                 on a large-scale corpus of single \
                 cell transcriptomes to enable context-aware \
                 predictions in settings with limited data in \
                 network biology.",
    packages=find_packages(),
    python_requires=">=3.10",
    include_package_data=True,
    install_requires=[
        "anndata",
        "bitsandbytes",
        "datasets",
        "loompy",
        "matplotlib",
        "numpy",
        "optuna",
        "optuna-integration",
        "packaging",
        "pandas",
        "peft",
        "pyarrow",
        "pytz",
        "ray",
        "scanpy",
        "scikit-learn",
        "scipy",
        "seaborn",
        "setuptools",
        "statsmodels",
        "tdigest",
        "tensorboard",
        "torch",
        "tqdm",
        "transformers",
    ],
)

"""阶段 perturb:in silico KO 扰动。

输入: workspace/tokenized/<prefix>.dataset
       + config.paths.model_dir(权重)
输出: workspace/isp_output/<prefix>_KO 相关 pickle(原始扰动结果)

通过独立 worker 子进程运行(避免主进程环境不干净),GPU 自适应默认值。
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from utils.config_loader import load_config  # noqa: E402
from utils.isp_runner import resolve_genes, run_worker  # noqa: E402


def _model_usable(model_dir: Path) -> bool:
    if not (model_dir / "config.json").exists():
        return False
    for name in ("model.safetensors", "pytorch_model.bin"):
        wf = model_dir / name
        if wf.exists() and wf.stat().st_size > 1_000_000:
            return True
    return False


def run_stage(cfg: dict, force: bool = False, log=None) -> dict:
    log = log or print
    work_dir = Path(cfg["project"]["work_dir"])
    name = cfg["project"]["name"]
    p_cfg = cfg["perturb"]

    # ---- 1. 模型检查 ----
    model_dir = Path(cfg["paths"]["model_dir"])
    if not _model_usable(model_dir):
        raise RuntimeError(
            f"模型权重不可用: {model_dir}\n"
            "请先运行 bash init.sh 下载权重,或确认 config.paths.model_dir 指向真实权重目录。"
        )
    log(f"[perturb] 使用模型: {model_dir}")

    # ---- 2. 输入 .dataset ----
    dataset_dir = work_dir / "tokenized"
    dataset_paths = sorted(dataset_dir.glob("*.dataset"))
    if not dataset_paths:
        raise FileNotFoundError(f"未找到 tokenized 产物: {dataset_dir} (请先跑 tokenize 阶段)")
    dataset_path = dataset_paths[0]
    log(f"[perturb] 输入 tokenized: {dataset_path}")

    # ---- 3. 解析扰动基因(symbol/Ensembl,空则自动挑 TF) ----
    import scanpy as sc

    prepared_h5ad = work_dir / "input_h5ad" / f"{name}_for_geneformer.h5ad"
    adata = sc.read_h5ad(prepared_h5ad) if prepared_h5ad.exists() else None
    ens_ids = resolve_genes(cfg, adata, log=log)
    if not ens_ids:
        raise ValueError("未解析到任何扰动基因,请检查 config.perturb.genes_to_perturb。")

    # ---- 4. 检查是否已有产物(缓存跳过) ----
    isp_output_dir = work_dir / "isp_output"
    isp_output_dir.mkdir(parents=True, exist_ok=True)
    isp_prefix = f"{name}_KO"
    existing = sorted(isp_output_dir.glob(f"{isp_prefix}*"))
    if existing and not force:
        log(f"[perturb] 已有扰动产物 {len(existing)} 个,跳过: {existing[:3]} ...")
        return {"output_directory": str(isp_output_dir), "output_prefix": isp_prefix}

    # ---- 5. 构造 payload 并调用 worker ----
    payload = {
        "geneformer_repo": str(cfg["paths"]["geneformer_repo"]),
        "model_directory": str(model_dir),
        "input_data_file": str(dataset_path),
        "output_directory": str(isp_output_dir),
        "output_prefix": isp_prefix,
        "perturb_type": p_cfg.get("perturb_type", "delete"),
        "genes_to_perturb": ens_ids,
        "combos": p_cfg.get("combos", 0),
        "anchor_gene": p_cfg.get("anchor_gene"),
        "model_type": "Pretrained",
        "num_classes": 0,
        "emb_mode": p_cfg.get("emb_mode", "cls"),
        "filter_data": None,
        "max_ncells": p_cfg.get("max_ncells"),
        "emb_layer": p_cfg.get("emb_layer", 0),
        "forward_batch_size": p_cfg.get("forward_batch_size"),
        "nproc": p_cfg.get("nproc"),
        "model_version": cfg["tokenize"].get("model_version", "V2"),
        "cuda_device": cfg["hpc"].get("cuda_device", "auto"),
    }
    result = run_worker("isp", payload, cfg, log=log)

    output_files = sorted(isp_output_dir.glob(f"{isp_prefix}*"))
    log(f"[perturb] 扰动完成,输出 {len(output_files)} 个文件到 {isp_output_dir}")
    return {
        "output_directory": str(isp_output_dir),
        "output_prefix": isp_prefix,
        "genes_perturbed": ens_ids,
        "worker_result": result,
    }


if __name__ == "__main__":
    cfg = load_config()
    print(run_stage(cfg, log=print))

"""阶段 data:加载 h5ad → 整理成 Geneformer 友好格式 → 抽样 → 保存中间 h5ad。

忠实复刻 AIVC_Geneformer_Norman2019_中文教程 的预处理逻辑:
  1. 使用原始 count 层(存在 'counts' 层则切过去)
  2. 补齐 Ensembl ID(var 无 ensembl_id 列时用 var_names)
  3. 计算 n_counts
  4. 保留 cell_type / perturbation_name / filter_pass 元信息
  5. fast_mode 时按固定种子抽样 max_cells 个细胞
产物: workspace/input_h5ad/<name>_for_geneformer.h5ad
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import scanpy as sc
from scipy import sparse


def _cache_marker(cfg: dict) -> Path:
    work_dir = Path(cfg["project"]["work_dir"])
    name = cfg["project"]["name"]
    out_dir = work_dir / "input_h5ad"
    return out_dir / f"{name}_for_geneformer.h5ad"


def run_stage(cfg: dict, force: bool = False, log=None) -> Path:
    log = log or print
    out_h5ad = _cache_marker(cfg)

    if out_h5ad.exists() and not force:
        log(f"[data] 缓存已存在,跳过: {out_h5ad} (如需重跑加 --force 或删除该文件)")
        return out_h5ad

    data_file = Path(cfg["paths"]["data_file"])
    if not data_file.exists():
        raise FileNotFoundError(
            f"数据文件不存在: {data_file}\n"
            "请先运行 bash init.sh 下载数据,或把 h5ad 放到 config.paths.data_file 指定位置。"
        )

    d_cfg = cfg["data"]
    cell_type_key = d_cfg.get("cell_type_key", "cell_type")
    perturbation_key = d_cfg.get("perturbation_key", "perturbation_name")

    log(f"[data] 读取: {data_file}")
    adata = sc.read_h5ad(data_file)

    # ---- 1. 使用原始 count 层 ----
    if "counts" in adata.layers:
        adata.X = adata.layers["counts"].copy()

    # ---- 2. 补齐 Ensembl ID ----
    if "ensembl_id" not in adata.var.columns:
        adata.var["ensembl_id"] = adata.var_names.astype(str)
    ens_ratio = float(np.mean(adata.var["ensembl_id"].astype(str).str.startswith("ENSG")))
    log(f"[data] Ensembl-like ID 比例: {ens_ratio:.3f}")
    if ens_ratio < 0.5:
        log(f"[data] 警告: 只有 {ens_ratio:.1%} 的基因是 ENSG 格式;"
            "Geneformer 词表按 Ensembl ID 匹配,请确认数据基因名是 Ensembl ID。")

    # ---- 3. 计算 n_counts ----
    if "n_counts" not in adata.obs.columns:
        if sparse.issparse(adata.X):
            adata.obs["n_counts"] = np.asarray(adata.X.sum(axis=1)).ravel()
        else:
            adata.obs["n_counts"] = adata.X.sum(axis=1)

    # ---- 4. 元信息兜底 ----
    if cell_type_key not in adata.obs.columns:
        adata.obs[cell_type_key] = "K562"
    if perturbation_key not in adata.obs.columns:
        adata.obs[perturbation_key] = "unknown"
    adata.obs[perturbation_key] = adata.obs[perturbation_key].astype(str)
    adata.obs["filter_pass"] = 1

    # ---- 可选过滤 ----
    if d_cfg.get("filter_cell_type"):
        keep = adata.obs[cell_type_key].astype(str) == str(d_cfg["filter_cell_type"])
        log(f"[data] 按 cell_type='{d_cfg['filter_cell_type']}' 过滤: {int(keep.sum())}/{adata.n_obs}")
        adata = adata[keep].copy()
    if d_cfg.get("filter_perturbation"):
        keep = adata.obs[perturbation_key].astype(str) == str(d_cfg["filter_perturbation"])
        log(f"[data] 按 perturbation='{d_cfg['filter_perturbation']}' 过滤: {int(keep.sum())}/{adata.n_obs}")
        adata = adata[keep].copy()

    # ---- 5. 快速模式抽样 ----
    if d_cfg.get("fast_mode", True) and adata.n_obs > int(d_cfg.get("max_cells", 5000)):
        max_cells = int(d_cfg.get("max_cells", 5000))
        seed = int(cfg["project"].get("seed", 123))
        log(f"[data] FAST_MODE: 从 {adata.n_obs} 细胞抽样 {max_cells} (seed={seed})")
        sc.pp.subsample(adata, n_obs=max_cells, random_state=seed)
    else:
        log(f"[data] 全量模式,共 {adata.n_obs} 细胞 / {adata.n_vars} 基因")

    # ---- 保存 ----
    out_h5ad.parent.mkdir(parents=True, exist_ok=True)
    adata.write_h5ad(out_h5ad)
    log(f"[data] 已保存中间 h5ad: {out_h5ad} ({adata.n_obs} 细胞 / {adata.n_vars} 基因)")
    return out_h5ad