| |
| """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 |
|
|
|
|
| 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() |
|
|
