scripts/run_dst_gnn.py

#!/usr/bin/env python
"""Run manuscript-aligned DST-GNN analysis on COSTE/SSS stage graphs."""

from __future__ import annotations

import argparse
import hashlib
import json
import platform
import random
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch

try:
    import matplotlib.pyplot as plt
except ModuleNotFoundError:  # pragma: no cover - optional runtime dependency
    plt = None

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from dst_gnn import (  # noqa: E402
    DSTGNN,
    STAGE_ORDER,
    build_temporal_graphs,
    explain_node_transition,
    rank_edges_by_stage_change,
    rank_nodes_by_embedding_drift,
    train_model,
)
from dst_gnn.explain import edge_mask_to_frame  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run DST-GNN on temporal COSTE/SSS graphs.")
    parser.add_argument("--csv", type=Path, required=True, help="Path to the flattened COSTE/SSS CSV.")
    parser.add_argument("--output-dir", type=Path, required=True, help="Directory for outputs.")
    parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducible training.")
    parser.add_argument("--hidden-channels", type=int, default=32, help="Hidden channel dimension.")
    parser.add_argument("--dropout", type=float, default=0.0, help="Dropout probability.")
    parser.add_argument("--epochs", type=int, default=400, help="Training epochs.")
    parser.add_argument("--lr", type=float, default=0.01, help="Adam learning rate.")
    parser.add_argument("--weight-decay", type=float, default=5e-4, help="Adam weight decay.")
    parser.add_argument("--top-k", type=int, default=20, help="Rows to keep in ranked outputs.")
    parser.add_argument(
        "--device",
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Training device, e.g. cpu or cuda.",
    )
    parser.add_argument(
        "--run-explainer",
        action="store_true",
        help="Run GNNExplainer on the highest-drift node for the T2 -> T3 transition.",
    )
    return parser.parse_args()


def set_reproducibility(seed: int) -> None:
    """Set deterministic seeds for the current process."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.use_deterministic_algorithms(True, warn_only=True)
    if hasattr(torch.backends, "cudnn"):
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def sha256_file(path: Path) -> str:
    """Return the SHA256 digest of a file."""

    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            chunk = handle.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def get_runtime_versions() -> dict[str, str | None]:
    """Collect runtime version information for reproducibility metadata."""

    versions: dict[str, str | None] = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "pandas": pd.__version__,
        "numpy": np.__version__,
        "matplotlib": getattr(plt, "__version__", None) if plt is not None else None,
    }
    try:
        import torch_geometric

        versions["torch_geometric"] = torch_geometric.__version__
    except ModuleNotFoundError:
        versions["torch_geometric"] = None
    return versions


def save_matrix(matrix: torch.Tensor | pd.DataFrame | np.ndarray, output_path: Path, node_names: list[str]) -> None:
    if isinstance(matrix, torch.Tensor):
        values = matrix.detach().cpu().numpy()
    elif isinstance(matrix, np.ndarray):
        values = matrix
    else:
        values = matrix.values
    frame = pd.DataFrame(values, index=node_names, columns=node_names)
    frame.to_csv(output_path)


def plot_training_curve(history: pd.DataFrame, output_path: Path) -> None:
    if plt is None:
        return
    plt.figure(figsize=(6, 4))
    plt.plot(history["epoch"], history["loss"], linewidth=2)
    plt.xlabel("Epoch")
    plt.ylabel("MSE loss")
    plt.title("DST-GNN training curve")
    plt.tight_layout()
    plt.savefig(output_path, dpi=200)
    plt.close()


def plot_top_nodes(nodes: pd.DataFrame, output_path: Path) -> None:
    if plt is None:
        return
    frame = nodes.iloc[:10].copy()
    frame = frame.iloc[::-1]
    plt.figure(figsize=(8, 5))
    plt.barh(frame["node"], frame["drift_score"])
    plt.xlabel("Embedding drift score")
    plt.title("Top dynamic cell types")
    plt.tight_layout()
    plt.savefig(output_path, dpi=200)
    plt.close()


def write_checksums(csv_path: Path, output_dir: Path) -> None:
    """Write SHA256 checksums for the input CSV and generated output files."""

    checksum_path = output_dir / "checksums.sha256"
    entries: list[tuple[str, str]] = [
        (sha256_file(csv_path), f"input_csv::{csv_path.name}"),
    ]
    for path in sorted(output_dir.iterdir()):
        if path.is_file() and path.name != checksum_path.name:
            entries.append((sha256_file(path), path.name))

    lines = [f"{digest}  {label}" for digest, label in entries]
    checksum_path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def build_run_config(
    args: argparse.Namespace,
    bundle,
    output_dir: Path,
    generated_files: list[str],
    csv_argument: str,
    output_argument: str,
) -> dict[str, object]:
    """Build the structured run configuration metadata."""

    return {
        "input_csv": {
            "path": csv_argument,
            "resolved_path": str(args.csv),
            "filename": args.csv.name,
            "sha256": sha256_file(args.csv),
        },
        "output_dir": output_argument,
        "resolved_output_dir": str(output_dir),
        "seed": args.seed,
        "device": args.device,
        "model": {
            "hidden_channels": args.hidden_channels,
            "dropout": args.dropout,
            "epochs": args.epochs,
            "lr": args.lr,
            "weight_decay": args.weight_decay,
            "top_k": args.top_k,
            "run_explainer": args.run_explainer,
        },
        "analysis": {
            "node_count": len(bundle.node_names),
            "stage_order": list(STAGE_ORDER),
            "sample_counts": bundle.sample_counts,
        },
        "environment": get_runtime_versions(),
        "generated_files": generated_files,
    }


def main() -> None:
    args = parse_args()
    csv_argument = str(args.csv)
    output_argument = str(args.output_dir)
    args.csv = args.csv.resolve()
    args.output_dir = args.output_dir.resolve()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    set_reproducibility(args.seed)

    bundle = build_temporal_graphs(args.csv)
    node_names = bundle.node_names
    graphs = [graph.to(args.device) for graph in bundle.stage_graphs]

    model = DSTGNN(
        num_nodes=len(node_names),
        hidden_channels=args.hidden_channels,
        dropout=args.dropout,
    ).to(args.device)

    history = pd.DataFrame(
        train_model(
            model,
            graphs,
            epochs=args.epochs,
            lr=args.lr,
            weight_decay=args.weight_decay,
        )
    )
    history.to_csv(args.output_dir / "training_history.csv", index=False)
    plot_training_curve(history, args.output_dir / "training_curve.png")

    model.eval()
    with torch.no_grad():
        embeddings = model.encode_observed_sequence(graphs)
        predictions = model.forecast_sequence(graphs)

    top_nodes = rank_nodes_by_embedding_drift(node_names, embeddings[0], embeddings[-1])
    top_nodes.head(args.top_k).to_csv(args.output_dir / "top_dynamic_nodes.csv", index=False)
    plot_top_nodes(top_nodes, args.output_dir / "top_dynamic_nodes.png")

    edge_changes = rank_edges_by_stage_change(
        node_names,
        bundle.stage_matrices["T1"],
        bundle.stage_matrices["T3"],
    )
    edge_changes.head(args.top_k).to_csv(args.output_dir / "top_edge_changes.csv", index=False)

    for stage in STAGE_ORDER:
        save_matrix(bundle.stage_matrices[stage], args.output_dir / f"observed_{stage}_sss.csv", node_names)

    for index, stage in enumerate(STAGE_ORDER[1:]):
        save_matrix(predictions[index], args.output_dir / f"predicted_{stage}_sss.csv", node_names)

    summary = {
        "csv_path": csv_argument,
        "resolved_csv_path": str(args.csv),
        "csv_sha256": sha256_file(args.csv),
        "device": args.device,
        "seed": args.seed,
        "node_count": len(node_names),
        "stage_order": list(STAGE_ORDER),
        "sample_counts": bundle.sample_counts,
        "top_dynamic_nodes": top_nodes.head(5)["node"].tolist(),
        "top_edge_change_pairs": edge_changes.head(5)[["source", "target"]].to_dict(orient="records"),
    }
    (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

    if args.run_explainer:
        target_node_name = str(top_nodes.iloc[0]["node"])
        target_node_index = node_names.index(target_node_name)
        explanation = explain_node_transition(
            model,
            graphs[1],
            target_node=target_node_index,
            prev_hidden=embeddings[0],
            epochs=200,
        )
        if getattr(explanation, "edge_mask", None) is not None:
            edge_frame = edge_mask_to_frame(graphs[1], node_names, explanation.edge_mask)
            edge_frame.head(args.top_k).to_csv(args.output_dir / "explainer_edges_t2_to_t3.csv", index=False)

    generated_files = sorted(
        path.name
        for path in args.output_dir.iterdir()
        if path.is_file()
    )
    for name in ("run_config.json", "checksums.sha256"):
        if name not in generated_files:
            generated_files.append(name)
    generated_files = sorted(generated_files)
    run_config = build_run_config(
        args,
        bundle,
        args.output_dir,
        generated_files,
        csv_argument,
        output_argument,
    )
    (args.output_dir / "run_config.json").write_text(
        json.dumps(run_config, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    write_checksums(args.csv, args.output_dir)

    print(f"Finished. Outputs written to {args.output_dir}")


if __name__ == "__main__":
    main()