path: root/kaft_mvp/experiment.py

"""End-to-end MVP experiment and artifact generation."""

from __future__ import annotations

import json
import platform
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch_geometric

from .data import load_cora
from .trainers import BPTrainer, KAFTTrainer, seed_everything


@dataclass(frozen=True)
class MVPConfig:
    dataset: str = "Cora"
    seeds: tuple[int, ...] = (0, 1, 2)
    depth: int = 6
    diagnostic_depth: int = 10
    hidden_dim: int = 64
    epochs: int = 200
    diagnostic_epochs: int = 100
    eval_every: int = 5
    learning_rate: float = 0.01
    weight_decay: float = 5e-4
    diffusion_alpha: float = 0.5
    diffusion_steps: int = 10
    hop_cap: int = 3
    num_probes: int = 64
    feedback_learning_rate: float = 0.5
    align_every: int = 10
    device: str = "cpu"


def _train(
    trainer: BPTrainer | KAFTTrainer,
    epochs: int,
    eval_every: int,
) -> tuple[dict[str, Any], dict[str, list[float | None]]]:
    history: dict[str, list[float | None]] = {
        "train_loss": [],
        "train_accuracy": [],
        "validation_accuracy": [],
        "test_accuracy": [],
    }
    best_validation = float("-inf")
    test_at_best_validation = float("nan")
    best_epoch = -1
    for epoch in range(epochs):
        step = trainer.train_step()
        history["train_loss"].append(step.loss)
        history["train_accuracy"].append(step.train_accuracy)
        if epoch % eval_every == 0 or epoch == epochs - 1:
            validation = trainer.accuracy("val_mask")
            test = trainer.accuracy("test_mask")
            history["validation_accuracy"].append(validation)
            history["test_accuracy"].append(test)
            if validation > best_validation:
                best_validation = validation
                test_at_best_validation = test
                best_epoch = epoch
        else:
            history["validation_accuracy"].append(None)
            history["test_accuracy"].append(None)
    return {
        "best_validation_accuracy": best_validation,
        "test_accuracy_at_validation_peak": test_at_best_validation,
        "best_epoch": best_epoch,
        "final_train_loss": history["train_loss"][-1],
        "final_train_accuracy": history["train_accuracy"][-1],
    }, history


def _make_trainer(
    method: str,
    data: dict[str, Any],
    config: MVPConfig,
    depth: int,
) -> BPTrainer | KAFTTrainer:
    common = {
        "data": data,
        "depth": depth,
        "hidden_dim": config.hidden_dim,
        "learning_rate": config.learning_rate,
        "weight_decay": config.weight_decay,
    }
    if method == "BP":
        return BPTrainer(**common)
    if method == "KAFT":
        return KAFTTrainer(
            **common,
            diffusion_alpha=config.diffusion_alpha,
            diffusion_steps=config.diffusion_steps,
            hop_cap=config.hop_cap,
            num_probes=config.num_probes,
            feedback_learning_rate=config.feedback_learning_rate,
            align_every=config.align_every,
        )
    raise ValueError(f"Unknown method: {method}")


def _summary_frame(records: list[dict[str, Any]]) -> pd.DataFrame:
    rows = []
    for method in ("BP", "KAFT"):
        values = np.asarray(
            [
                row["result"]["test_accuracy_at_validation_peak"]
                for row in records
                if row["method"] == method
            ],
            dtype=float,
        )
        rows.append(
            {
                "method": method,
                "mean_test_accuracy_percent": 100.0 * values.mean(),
                "std_test_accuracy_percent": (
                    100.0 * values.std(ddof=1)
                    if len(values) > 1
                    else 0.0
                ),
                "seeds": len(values),
            }
        )
    return pd.DataFrame(rows)


def _plot_histories(
    records: list[dict[str, Any]],
    output_path: Path,
) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
    for method, color in (("BP", "#4C78A8"), ("KAFT", "#E45756")):
        histories = [
            row["history"]
            for row in records
            if row["method"] == method
        ]
        loss = np.asarray([history["train_loss"] for history in histories])
        test = np.asarray(
            [
                [
                    np.nan if value is None else value
                    for value in history["test_accuracy"]
                ]
                for history in histories
            ]
        )
        epochs = np.arange(1, loss.shape[1] + 1)
        median_loss = np.nanmedian(loss, axis=0)
        axes[0].plot(epochs, median_loss, color=color, label=method)
        axes[0].fill_between(
            epochs,
            np.nanpercentile(loss, 25, axis=0),
            np.nanpercentile(loss, 75, axis=0),
            color=color,
            alpha=0.15,
        )
        evaluated = ~np.isnan(test).all(axis=0)
        axes[1].plot(
            epochs[evaluated],
            100.0 * np.nanmedian(test[:, evaluated], axis=0),
            color=color,
            label=method,
        )
    axes[0].set_xlabel("Epoch")
    axes[0].set_ylabel("Training cross-entropy")
    axes[0].set_yscale("log")
    axes[1].set_xlabel("Epoch")
    axes[1].set_ylabel("Test accuracy (%)")
    for axis in axes:
        axis.grid(alpha=0.2)
        axis.legend(frameon=False)
    fig.suptitle("Cora, identical 6-layer GCN forward model")
    fig.tight_layout()
    fig.savefig(output_path, dpi=180, bbox_inches="tight")
    plt.close(fig)


def run_mvp(
    config: MVPConfig | None = None,
    output_dir: str | Path = "artifacts",
    data_root: str | Path = "data",
) -> dict[str, Any]:
    """Run BP/KAFT accuracy and BP transport diagnostics end to end."""
    config = config or MVPConfig()
    if config.dataset != "Cora":
        raise ValueError("The minimal package currently exposes only Cora")
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    data = load_cora(root=data_root, device=config.device)

    records: list[dict[str, Any]] = []
    for seed in config.seeds:
        for method in ("BP", "KAFT"):
            seed_everything(seed)
            trainer = _make_trainer(method, data, config, config.depth)
            result, history = _train(
                trainer,
                config.epochs,
                config.eval_every,
            )
            records.append(
                {
                    "seed": seed,
                    "method": method,
                    "result": result,
                    "history": history,
                }
            )

    diagnostics: list[dict[str, Any]] = []
    for seed in config.seeds:
        seed_everything(seed)
        trainer = _make_trainer(
            "BP",
            data,
            config,
            config.diagnostic_depth,
        )
        result, _ = _train(
            trainer,
            config.diagnostic_epochs,
            config.eval_every,
        )
        diagnostics.append(
            {
                "seed": seed,
                "training_result": result,
                **trainer.gradient_diagnostic(),
            }
        )

    summary = _summary_frame(records)
    summary.to_csv(output_dir / "mvp_summary.csv", index=False)
    _plot_histories(records, output_dir / "training_curves.png")

    payload = {
        "config": {
            **asdict(config),
            "seeds": list(config.seeds),
        },
        "environment": {
            "python": platform.python_version(),
            "torch": torch.__version__,
            "torch_geometric": torch_geometric.__version__,
            "device": config.device,
        },
        "summary": summary.to_dict(orient="records"),
        "records": records,
        "gradient_diagnostics": diagnostics,
    }
    with (output_dir / "mvp_results.json").open("w") as handle:
        json.dump(payload, handle, indent=2, allow_nan=False)
    return payload