train.py

from __future__ import annotations

import json
import math
import random
import shutil
from pathlib import Path

import numpy as np
import torch
import trackio
from model import MeshGraphGAT, parameter_count
from safetensors.torch import save_file
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from torch.nn import functional as F

PROJECT_DIR = Path(__file__).resolve().parent
ROOT_DIR = PROJECT_DIR.parents[1]
SOURCE_DIR = ROOT_DIR / "projects" / "meshgraph-gcn"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "meshgraph-gat"
DATA_DIR = PROJECT_DIR / "data"
SEED = 2189


def metrics(labels: np.ndarray, scores: np.ndarray, threshold: float) -> dict:
    predictions = scores >= threshold
    return {
        "accuracy": float(accuracy_score(labels, predictions)),
        "roc_auc": float(roc_auc_score(labels, scores)),
        "average_precision": float(average_precision_score(labels, scores)),
        "precision": float(precision_score(labels, predictions, zero_division=0)),
        "recall": float(recall_score(labels, predictions, zero_division=0)),
        "f1": float(f1_score(labels, predictions, zero_division=0)),
    }


def best_threshold(labels: np.ndarray, scores: np.ndarray) -> float:
    candidates = np.quantile(scores, np.linspace(0.05, 0.95, 300))
    return float(
        max(
            candidates,
            key=lambda threshold: f1_score(
                labels, scores >= threshold, zero_division=0
            ),
        )
    )


def attention_audit(
    attention: torch.Tensor,
    adjacency: torch.Tensor,
    nodes: np.ndarray,
) -> dict:
    mean_attention = attention.mean(0)
    normalized_entropies = []
    top_edges = []
    for node in nodes:
        neighbors = torch.nonzero(adjacency[node] > 0).flatten()
        weights = mean_attention[node, neighbors]
        if len(neighbors) > 1:
            entropy = -(weights * weights.clamp_min(1e-12).log()).sum()
            normalized_entropies.append(float(entropy / math.log(len(neighbors))))
        for neighbor, weight in zip(neighbors.tolist(), weights.tolist(), strict=True):
            top_edges.append((float(weight), int(node), int(neighbor)))
    top_edges.sort(reverse=True)
    return {
        "mean_normalized_attention_entropy": float(np.mean(normalized_entropies)),
        "top_test_edges": [
            {"source": source, "target": target, "attention": weight}
            for weight, source, target in top_edges[:20]
        ],
    }


def main() -> None:
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    torch.set_num_threads(1)
    graph = np.load(SOURCE_DIR / "data" / "meshgraph.npz")
    features = graph["features"].astype(np.float32)
    labels = graph["labels"].astype(np.int64)
    train_indices = graph["train_indices"]
    validation_indices = graph["validation_indices"]
    test_indices = graph["test_indices"]
    mean = features[train_indices].mean(0)
    scale = np.maximum(features[train_indices].std(0), 1e-5)
    feature_tensor = torch.from_numpy((features - mean) / scale)
    label_tensor = torch.from_numpy(labels)
    adjacency = torch.from_numpy(graph["adjacency"].astype(np.float32))
    model = MeshGraphGAT(features.shape[1])
    positive_weight = (labels[train_indices] == 0).sum() / max(
        1, (labels[train_indices] == 1).sum()
    )
    class_weights = torch.tensor([1.0, positive_weight], dtype=torch.float32)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.008, weight_decay=0.002)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=400)
    best_auc = -1.0
    best_epoch = 0
    best_state = None
    trackio.init(
        project="meshgraph-gat",
        name="four-head-enterprise-gat-v1",
        config={
            "parameters": parameter_count(model),
            "nodes": len(labels),
            "edges": int(adjacency.sum() // 2),
            "heads": 4,
        },
    )
    for epoch in range(1, 401):
        model.train()
        logits = model(feature_tensor, adjacency)
        loss = F.cross_entropy(
            logits[train_indices], label_tensor[train_indices], weight=class_weights
        )
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        scheduler.step()
        model.eval()
        with torch.inference_mode():
            scores = model(feature_tensor, adjacency).softmax(1)[:, 1]
        validation_auc = roc_auc_score(
            labels[validation_indices], scores[validation_indices].numpy()
        )
        if validation_auc > best_auc:
            best_auc = validation_auc
            best_epoch = epoch
            best_state = {
                name: value.detach().cpu().clone()
                for name, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 20 == 0:
            trackio.log(
                {
                    "epoch": epoch,
                    "training_loss": float(loss.detach()),
                    "validation_roc_auc": validation_auc,
                }
            )
    assert best_state is not None
    model.load_state_dict(best_state)
    model.eval()
    with torch.inference_mode():
        logits, attention = model(
            feature_tensor, adjacency, return_attention=True
        )
        scores = logits.softmax(1)[:, 1].numpy()
    threshold = best_threshold(labels[validation_indices], scores[validation_indices])
    gcn_report = json.loads(
        (
            SOURCE_DIR / "artifacts" / "meshgraph-gcn" / "evaluation.json"
        ).read_text(encoding="utf-8")
    )
    results = {
        "model": "MeshGraph GAT",
        "parameters": parameter_count(model),
        "best_epoch": best_epoch,
        "best_validation_roc_auc": best_auc,
        "threshold": threshold,
        "gat_test": metrics(labels[test_indices], scores[test_indices], threshold),
        "existing_gcn_test": gcn_report["gcn_test"],
        "existing_feature_only_test": gcn_report["feature_only_logistic_test"],
        "attention_audit": attention_audit(attention, adjacency, test_indices),
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
    np.savez(
        ARTIFACT_DIR / "preprocessing.npz",
        mean=mean,
        scale=scale,
        threshold=threshold,
    )
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2), encoding="utf-8"
    )
    shutil.copy2(SOURCE_DIR / "data" / "meshgraph.npz", DATA_DIR / "meshgraph.npz")
    trackio.log(
        {
            "test_roc_auc": results["gat_test"]["roc_auc"],
            "test_average_precision": results["gat_test"]["average_precision"],
            "test_f1": results["gat_test"]["f1"],
            "attention_entropy": results["attention_audit"][
                "mean_normalized_attention_entropy"
            ],
        }
    )
    trackio.finish()
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()