scripts/inference.py

#!/usr/bin/env python3
"""Run MVGNN-PPIS five-fold ensemble inference."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import pandas as pd
import torch


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from model import MVGNN, Seed_everything, model_test  # noqa: E402


def project_path(value: str | Path) -> Path:
    """Resolve package-relative paths while preserving explicit absolute paths."""
    path = Path(value).expanduser()
    return path if path.is_absolute() else PROJECT_ROOT / path


def load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Predict protein-protein interaction sites with MVGNN-PPIS."
    )
    parser.add_argument(
        "--config",
        default="conf/config.json",
        help="Package-relative JSON configuration file.",
    )
    parser.add_argument("--dataset", help="Override the evaluation CSV path.")
    parser.add_argument("--feature-path", help="Override the precomputed feature directory.")
    parser.add_argument("--weight-path", help="Override the checkpoint directory.")
    parser.add_argument("--output-path", help="Override the prediction output directory.")
    parser.add_argument("--num-workers", type=int, help="DataLoader worker count.")
    parser.add_argument("--seed", type=int, help="Random seed.")
    parser.add_argument(
        "--device",
        help="PyTorch device, for example 'cuda', 'cuda:0', or 'cpu'.",
    )
    return parser.parse_args()


def required_feature_files(feature_dir: Path, protein_ids: list[str]) -> list[Path]:
    suffixes = ("X", "node_feature", "mask", "label", "adj")
    return [
        feature_dir / f"{protein_id}_{suffix}.tensor"
        for protein_id in protein_ids
        for suffix in suffixes
    ]


def main() -> None:
    args = parse_args()
    config = load_config(project_path(args.config))

    paths = config["paths"]
    dataset_path = project_path(args.dataset or paths["dataset"])
    feature_dir = project_path(args.feature_path or paths["feature_dir"])
    weight_dir = project_path(args.weight_path or paths["weight_dir"])
    output_dir = project_path(args.output_path or paths["output_dir"])

    runtime = config["runtime"]
    seed = args.seed if args.seed is not None else runtime["seed"]
    num_workers = (
        args.num_workers if args.num_workers is not None else runtime["num_workers"]
    )
    requested_device = args.device or runtime.get("device", "auto")
    if requested_device == "auto":
        requested_device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(requested_device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("A CUDA/HIP device was requested but torch.cuda.is_available() is False.")

    test_df = pd.read_csv(
        dataset_path,
        dtype={"ID": str, "sequence": str, "label": str},
        keep_default_na=False,
    )
    required_columns = {"ID", "sequence", "label"}
    missing_columns = required_columns.difference(test_df.columns)
    if missing_columns:
        raise ValueError(
            f"Dataset {dataset_path} is missing columns: {sorted(missing_columns)}"
        )

    test_df["ID"] = test_df["ID"].str.strip()
    protein_ids = test_df["ID"].str.strip().drop_duplicates().tolist()
    missing_features = [
        path for path in required_feature_files(feature_dir, protein_ids) if not path.is_file()
    ]
    if missing_features:
        preview = "\n".join(f"  - {path}" for path in missing_features[:10])
        raise FileNotFoundError(
            f"Missing {len(missing_features)} precomputed feature files.\n{preview}"
        )

    folds = config["model"]["folds"]
    missing_weights = [
        weight_dir / f"fold{fold}.ckpt"
        for fold in range(folds)
        if not (weight_dir / f"fold{fold}.ckpt").is_file()
    ]
    if missing_weights:
        formatted = "\n".join(f"  - {path}" for path in missing_weights)
        raise FileNotFoundError(f"Missing model checkpoints:\n{formatted}")

    all_protein_data = {}
    for protein_id in protein_ids:
        all_protein_data[protein_id] = tuple(
            torch.load(
                feature_dir / f"{protein_id}_{suffix}.tensor",
                map_location="cpu",
                weights_only=True,
            )
            for suffix in ("X", "node_feature", "mask", "label", "adj")
        )

    Seed_everything(seed)
    model_config = dict(config["model"])
    model_config.update(
        {
            "seed": seed,
            "id_name": "ID",
            "remark": "PRO binding site prediction",
        }
    )
    print(f"Project root: {PROJECT_ROOT}")
    print(f"Dataset: {dataset_path}")
    print(f"Proteins: {len(test_df)}")
    print(f"Device: {device}")
    print(f"Checkpoints: {folds}")

    model_test(
        test_df,
        all_protein_data,
        MVGNN,
        model_config,
        weight_dir=weight_dir,
        output_dir=output_dir,
        logit=True,
        device=device,
        num_workers=num_workers,
    )


if __name__ == "__main__":
    main()