15_07_expt/scripts/model_gcn.py

"""Spatial-GCN stem -> Transformer-encoder architecture for figure-skating action classification.

Two temporal-modeling stages:
  1. Spatial GCN stem: treats the 17 COCO joints as a graph (edges = anatomical bones from
     preprocessing.BONE_PAIRS), applies Kipf-Welling style graph convolution per frame
     (H' = norm(A) @ H @ W), independently at every timestep -- no temporal mixing yet.
  2. Transformer encoder: sinusoidal positional encoding + a stack of full self-attn+FFN
     encoder layers over the per-frame graph-pooled representations (the temporal stage).
Followed by the same temporal mean-pool + deep GELU dense head used throughout this project,
so results are directly comparable to the conv-backbone and BiLSTM variants.

The GCN stem's node input is controlled by --node-features:
  "coords" (default): just the (17, 2) joint positions -- preprocessing.build_feature_tensor's
      first 34 columns, reshaped back into a graph (see preprocessing.py:284-288).
  "full": all 94 columns, scattered onto their anatomically natural node slot instead of being
      discarded or dumped in as an undifferentiated vector:
        - coords (34)            -> (x, y) at every joint
        - velocities (16)        -> (vx, vy) at the 8 KEY_JOINTS, zero elsewhere
        - angles (12)            -> the angle's *vertex* joint (JOINT_ANGLE_TRIPLETS[t][1]);
                                     a joint can be the vertex of up to 2 triplets, zero-padded
        - angular velocities (12)-> same vertex/slot pairing as angles
        - bone vectors (20)      -> BOTH endpoints of the bone (BONE_PAIRS is already the graph's
                                     edge list), signed +1 at the parent / -1 at the child so it
                                     reads as "which way my neighboring bone points relative to me";
                                     max degree in this skeleton is 2, zero-padded
  This uses every one of the 94 input dims, placed where it has a real anatomical meaning, without
  rewriting the conv into a full edge-conditioned message-passing layer -- bone vectors just ride
  along as extra node channels at both of their endpoints instead of a separate edge computation.
  All index/sign mappings are derived programmatically from preprocessing.BONE_PAIRS /
  JOINT_ANGLE_TRIPLETS / KEY_JOINTS, not hand-typed, so they can't silently drift out of sync.

Usage:
    python model_gcn.py --data-dir /path/to/processed
    python model_gcn.py --data-dir /path/to/processed --coarse --layers 6
    python model_gcn.py --data-dir /path/to/processed --coarse --node-features full
    python model_gcn.py --smoke
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import f1_score
from sklearn.utils.class_weight import compute_class_weight
from torch.utils.data import DataLoader, TensorDataset

try:
    from . import labels as labels_mod
    from . import preprocessing
except ImportError:
    import labels as labels_mod
    import preprocessing

from model import (
    DenseBlock,
    load_split,
    resolve_num_classes,
    assert_label_consistency,
    report_metrics,
    run_epoch,
    predict,
)
from model_transformers import PositionalEncoding, make_transformer_encoder

# ---------------------------------------------------------------------------
# Config (identical to model.py / model_transformer_bilstm.py for comparability)
# ---------------------------------------------------------------------------

EPOCHS = 100
BATCH_SIZE = 64
LEARNING_RATE = 5e-4
WEIGHT_DECAY = 1e-4
GRAD_CLIP = 1.0
MAX_CLASS_WEIGHT = 5.0
EARLY_STOP_PATIENCE = 20
SEED = 42
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

D_MODEL = 384
N_ATTN_BLOCKS = 3
GCN_CHANNELS = (64, 128)          # first GCN block's input channels vary by --node-features
NUM_JOINTS = preprocessing.NUM_JOINTS  # 17, COCO layout
COORD_DIMS = 2 * NUM_JOINTS       # 34: the flattened-coords prefix of the 94-dim feature vector

# Column offsets of each block within the 94-dim feature vector, matching
# preprocessing.build_feature_tensor's concatenation order: [coords(34), angles(12),
# bone_vectors(20), velocities(16), angular_velocities(12)].
_COORD_OFFSET = 0
_ANGLE_OFFSET = COORD_DIMS
_BONE_OFFSET = _ANGLE_OFFSET + len(preprocessing.JOINT_ANGLE_TRIPLETS)
_VEL_OFFSET = _BONE_OFFSET + 2 * len(preprocessing.BONE_PAIRS)
_ANGVEL_OFFSET = _VEL_OFFSET + 2 * len(preprocessing.KEY_JOINTS)
_TOTAL_DIMS = _ANGVEL_OFFSET + len(preprocessing.JOINT_ANGLE_TRIPLETS)
assert _TOTAL_DIMS == 94, f"offset bookkeeping drifted from preprocessing.py: got {_TOTAL_DIMS}"


def build_full_node_layout(num_joints: int = NUM_JOINTS):
    """Derive the (17, C) gather-index + sign-multiplier tensors that scatter all 94 feature
    columns onto their anatomically natural node slot. See module docstring for the layout.
    Computed from preprocessing.py's constants, not hand-typed, so it can't drift out of sync.
    """
    triplets = preprocessing.JOINT_ANGLE_TRIPLETS
    bones = preprocessing.BONE_PAIRS
    key_joints = preprocessing.KEY_JOINTS

    vertex_triplets = [[] for _ in range(num_joints)]
    for t_idx, (_, vertex, _c) in enumerate(triplets):
        vertex_triplets[vertex].append(t_idx)
    max_angle_slots = max(len(v) for v in vertex_triplets)

    incident_bones = [[] for _ in range(num_joints)]
    for e_idx, (parent, child) in enumerate(bones):
        incident_bones[parent].append((e_idx, 1.0))
        incident_bones[child].append((e_idx, -1.0))
    max_bone_slots = max(len(v) for v in incident_bones)

    key_joint_pos = {j: k for k, j in enumerate(key_joints)}

    channels = 2 + 2 + max_angle_slots + max_angle_slots + 2 * max_bone_slots
    src_idx = np.zeros((num_joints, channels), dtype=np.int64)
    mult = np.zeros((num_joints, channels), dtype=np.float32)

    for j in range(num_joints):
        # (x, y) -- always valid
        src_idx[j, 0] = _COORD_OFFSET + 2 * j
        src_idx[j, 1] = _COORD_OFFSET + 2 * j + 1
        mult[j, 0] = mult[j, 1] = 1.0

        # (vx, vy) -- only for KEY_JOINTS
        if j in key_joint_pos:
            k = key_joint_pos[j]
            src_idx[j, 2] = _VEL_OFFSET + 2 * k
            src_idx[j, 3] = _VEL_OFFSET + 2 * k + 1
            mult[j, 2] = mult[j, 3] = 1.0

        # angle slots, zero-padded up to max_angle_slots
        angle_col = 4
        for slot in range(max_angle_slots):
            if slot < len(vertex_triplets[j]):
                t = vertex_triplets[j][slot]
                src_idx[j, angle_col + slot] = _ANGLE_OFFSET + t
                mult[j, angle_col + slot] = 1.0

        # angular-velocity slots, same triplet pairing as the angle slots
        angvel_col = angle_col + max_angle_slots
        for slot in range(max_angle_slots):
            if slot < len(vertex_triplets[j]):
                t = vertex_triplets[j][slot]
                src_idx[j, angvel_col + slot] = _ANGVEL_OFFSET + t
                mult[j, angvel_col + slot] = 1.0

        # bone-vector slots (x, y pairs), signed +1 at parent / -1 at child, zero-padded
        bone_col = angvel_col + max_angle_slots
        for slot in range(max_bone_slots):
            if slot < len(incident_bones[j]):
                e, sign = incident_bones[j][slot]
                src_idx[j, bone_col + 2 * slot] = _BONE_OFFSET + 2 * e
                src_idx[j, bone_col + 2 * slot + 1] = _BONE_OFFSET + 2 * e + 1
                mult[j, bone_col + 2 * slot] = sign
                mult[j, bone_col + 2 * slot + 1] = sign

    return torch.from_numpy(src_idx), torch.from_numpy(mult), channels


FULL_NODE_SRC_IDX, FULL_NODE_MULT, FULL_NODE_CHANNELS = build_full_node_layout()


def make_gcn_channel_schedule(n_layers: int, start: int = 64, end: int = 128) -> tuple[int, ...]:
    """Channel width per GCN block, ramping start->end over n_layers (rounded to a multiple of 8).

    n_layers=2 reproduces the original fixed (64, 128) schedule exactly, so old checkpoints /
    default behavior are unaffected.
    """
    if n_layers == 1:
        return (end,)
    raw = np.linspace(start, end, n_layers)
    return tuple(int(round(c / 8) * 8) for c in raw)


# ---------------------------------------------------------------------------
# Graph adjacency (fixed skeleton topology, from preprocessing.BONE_PAIRS)
# ---------------------------------------------------------------------------


def build_normalized_adjacency(bone_pairs, num_joints: int = NUM_JOINTS) -> np.ndarray:
    """Symmetric-normalized adjacency D^-1/2 (A+I) D^-1/2 (Kipf & Welling).

    A+I includes self-loops so each node's own features survive aggregation, not just its
    neighbors'. Edges come from the anatomical bone list already used for bone-vector features,
    so the graph structure matches what the rest of the pipeline already treats as "connected".
    """
    A = np.eye(num_joints, dtype=np.float64)
    for a, b in bone_pairs:
        A[a, b] = 1.0
        A[b, a] = 1.0
    deg = A.sum(axis=1)
    deg_inv_sqrt = np.power(deg, -0.5)
    D_inv_sqrt = np.diag(deg_inv_sqrt)
    return D_inv_sqrt @ A @ D_inv_sqrt


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------


class GCNBlock(nn.Module):
    """Graph conv + BN + Dropout with a (projected) residual connection.

    Input/return: (B, T, J, Cin/Cout). Mirrors model.py's ConvBlock, just with graph
    aggregation (matmul against the fixed normalized adjacency) in place of a 1D conv.
    """

    def __init__(self, in_ch: int, out_ch: int, dropout: float = 0.15):
        super().__init__()
        self.shortcut = nn.Linear(in_ch, out_ch) if in_ch != out_ch else nn.Identity()
        self.lin = nn.Linear(in_ch, out_ch)
        self.bn = nn.BatchNorm1d(out_ch)
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
        agg = torch.matmul(A, x)              # (B, T, J, Cin): each node <- neighbors + self
        y = F.relu(self.lin(agg))              # (B, T, J, Cout)
        B, T, J, C = y.shape
        y = self.drop(self.bn(y.reshape(B * T * J, C)).reshape(B, T, J, C))
        return self.shortcut(x) + y


class SkatingGCNClassifier(nn.Module):
    """(B, T, 94) standardized feature sequence -> (B, num_classes) class logits.

    Stage 1 (spatial): reconstruct (B, T, 17, 2) joint coords from the feature vector's
    coordinate prefix, run a stack of residual GCN blocks per frame (no temporal mixing).
    Stage 2 (temporal): flatten per-frame joint features, project to D_MODEL, positional
    encoding + Transformer encoder (self-attn + FFN per layer) across time.
    Then temporal mean pool -> the same dense head used everywhere else in this project.
    """

    def __init__(self, num_classes: int, num_layers: int = N_ATTN_BLOCKS,
                 gcn_channels: tuple[int, ...] = GCN_CHANNELS, node_features: str = "coords"):
        super().__init__()
        if node_features not in ("coords", "full"):
            raise ValueError(f"node_features must be 'coords' or 'full', got {node_features!r}")
        self.node_features = node_features
        A = build_normalized_adjacency(preprocessing.BONE_PAIRS, NUM_JOINTS)
        self.register_buffer("A", torch.tensor(A, dtype=torch.float32))

        if node_features == "full":
            self.register_buffer("node_src_idx", FULL_NODE_SRC_IDX)
            self.register_buffer("node_mult", FULL_NODE_MULT)
            in_ch = FULL_NODE_CHANNELS
        else:
            in_ch = 2

        chans = (in_ch,) + tuple(gcn_channels)
        self.gcn_blocks = nn.ModuleList(
            [GCNBlock(chans[i], chans[i + 1]) for i in range(len(chans) - 1)]
        )
        self.proj = nn.Linear(NUM_JOINTS * gcn_channels[-1], D_MODEL)
        self.pos_enc = PositionalEncoding(D_MODEL)
        self.encoder = make_transformer_encoder(dim=D_MODEL, heads=8, num_layers=num_layers)
        self.head = nn.Sequential(
            DenseBlock(D_MODEL, 1024, 0.5),
            DenseBlock(1024, 512, 0.4),
            DenseBlock(512, 256, 0.3),
            nn.LayerNorm(256),
        )
        self.out = nn.Linear(256, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        if self.node_features == "full":
            h = x[..., self.node_src_idx] * self.node_mult   # (B, T, 17, FULL_NODE_CHANNELS)
        else:
            h = x[..., :COORD_DIMS].reshape(B, T, NUM_JOINTS, 2)   # (B, T, 17, 2)
        for blk in self.gcn_blocks:
            h = blk(h, self.A)
        h = h.reshape(B, T, -1)                # (B, T, 17*Cout)
        h = self.proj(h)                        # (B, T, D_MODEL)
        h = self.pos_enc(h)
        h = self.encoder(h)
        h = h.mean(dim=1)                        # temporal average pool -> (B, D_MODEL)
        return self.out(self.head(h))


# ---------------------------------------------------------------------------
# Train / evaluate (mirrors model_transformer_bilstm.py's train())
# ---------------------------------------------------------------------------


def train(data_dir: Path, coarse: bool = False, num_layers: int = N_ATTN_BLOCKS,
          node_features: str = "coords", gcn_layers: int = len(GCN_CHANNELS),
          noise_std: float = 0.0) -> dict:
    torch.manual_seed(SEED)
    np.random.seed(SEED)

    fine_n = resolve_num_classes(data_dir)
    taxonomy = labels_mod.COARSE_TAXONOMY if coarse else labels_mod.TAXONOMY
    num_classes = len(taxonomy)
    Xtr, ytr = load_split(data_dir, "train")
    Xva, yva = load_split(data_dir, "val")
    Xte, yte = load_split(data_dir, "test")

    assert_label_consistency(fine_n, ytr, yva, yte)
    if coarse:
        def coarsen(y):
            return np.array([labels_mod.FINE_TO_COARSE_IDX[int(v)] for v in y], dtype=np.int64)
        ytr, yva, yte = coarsen(ytr), coarsen(yva), coarsen(yte)
        assert_label_consistency(num_classes, ytr, yva, yte)
    in_features = Xtr.shape[-1]
    assert in_features >= COORD_DIMS, (
        f"GCN stem needs the flattened-coords prefix (first {COORD_DIMS} of 94 features) but "
        f"this data has only {in_features} feature columns -- wrong/incompatible processed dir."
    )
    print(f"[GCN] label space: {'COARSE action-level' if coarse else 'FINE'} | {num_classes} classes")

    mu = Xtr.mean(axis=(0, 1), keepdims=True)
    sd = Xtr.std(axis=(0, 1), keepdims=True) + 1e-6
    Xtr, Xva, Xte = (Xtr - mu) / sd, (Xva - mu) / sd, (Xte - mu) / sd

    print(f"[GCN] device={DEVICE} | num_classes={num_classes} | in_features={in_features} | "
          f"noise_std={noise_std} (train-only Gaussian input noise, val/test always clean)")
    print(f"[GCN] shapes: train={Xtr.shape} val={Xva.shape} test={Xte.shape}")
    print(f"[GCN] train classes present: {sorted(set(ytr.tolist()))}")

    present = np.unique(ytr)
    cw = np.clip(compute_class_weight(class_weight="balanced", classes=present, y=ytr),
                 None, MAX_CLASS_WEIGHT)
    weight = torch.ones(num_classes)
    for c, w in zip(present, cw):
        weight[int(c)] = float(w)
    criterion = nn.CrossEntropyLoss(weight=weight.to(DEVICE))

    gcn_channels = make_gcn_channel_schedule(gcn_layers)
    model = SkatingGCNClassifier(num_classes, num_layers=num_layers, gcn_channels=gcn_channels,
                                  node_features=node_features).to(DEVICE)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[GCN] model params: {n_params/1e6:.2f}M | encoder layers: {num_layers} | "
          f"gcn layers: {gcn_layers} {gcn_channels} | node_features={node_features}")
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)

    tr_loader = DataLoader(
        TensorDataset(torch.from_numpy(Xtr), torch.from_numpy(ytr)),
        batch_size=BATCH_SIZE, shuffle=True, drop_last=False,
    )
    va_loader = DataLoader(
        TensorDataset(torch.from_numpy(Xva), torch.from_numpy(yva)),
        batch_size=BATCH_SIZE, shuffle=False,
    )

    best_val_f1, best_state, since_improved = -1.0, None, 0
    for epoch in range(1, EPOCHS + 1):
        tr_loss, tr_acc = run_epoch(model, tr_loader, criterion, optimizer, noise_std=noise_std)
        va_loss, va_acc = run_epoch(model, va_loader, criterion)
        va_f1 = f1_score(yva, predict(model, Xva), labels=list(range(num_classes)),
                         average="macro", zero_division=0)
        if va_f1 > best_val_f1:
            best_val_f1, since_improved = va_f1, 0
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        else:
            since_improved += 1
        if epoch % 5 == 0 or epoch == 1:
            print(f"[GCN] epoch {epoch:3d} | train loss {tr_loss:.3f} acc {tr_acc:.3f} "
                  f"| val loss {va_loss:.3f} acc {va_acc:.3f} f1(macro) {va_f1:.3f}")
        if since_improved >= EARLY_STOP_PATIENCE:
            print(f"[GCN] early stop at epoch {epoch} (no val-F1 improvement for {EARLY_STOP_PATIENCE})")
            break

    if best_state is not None:
        model.load_state_dict(best_state)
    print(f"\n[GCN] restored best model (val macro-F1 = {best_val_f1:.3f})")

    metrics = report_metrics(yte, predict(model, Xte), taxonomy)
    layer_suffix = "" if num_layers == N_ATTN_BLOCKS else f"_{num_layers}L"
    feat_suffix = "" if node_features == "coords" else f"_{node_features}feat"
    gcn_suffix = "" if gcn_layers == len(GCN_CHANNELS) else f"_gcn{gcn_layers}L"
    noise_suffix = "" if noise_std == 0.0 else f"_noise{noise_std}"
    ckpt_name = f"model_gcn{'_coarse' if coarse else ''}{layer_suffix}{feat_suffix}{gcn_suffix}{noise_suffix}.pt"
    torch.save({"state_dict": model.state_dict(), "num_classes": num_classes,
                "in_features": in_features, "taxonomy": taxonomy, "coarse": coarse,
                "num_layers": num_layers, "node_features": node_features, "gcn_layers": gcn_layers,
                "gcn_channels": gcn_channels, "noise_std": noise_std,
                "feature_mean": mu, "feature_std": sd},
               data_dir / ckpt_name)
    print(f"\n[GCN] saved model -> {data_dir / ckpt_name}")
    return metrics


def smoke() -> None:
    """Validate the full plumbing on synthetic data shaped exactly like the pipeline output."""
    print("[GCN] SMOKE: synthetic (N,64,94) data, labels in TAXONOMY index space")
    num_classes = len(labels_mod.TAXONOMY)
    rng = np.random.default_rng(0)
    X = rng.standard_normal((40, 64, 94)).astype(np.float32)
    y = rng.integers(0, num_classes, size=40).astype(np.int64)
    assert_label_consistency(num_classes, y)

    print(f"[GCN] full-feature node layout: {FULL_NODE_CHANNELS} channels/joint "
          f"(src_idx range [{FULL_NODE_SRC_IDX.min()},{FULL_NODE_SRC_IDX.max()}], "
          f"nonzero mult entries: {int((FULL_NODE_MULT != 0).sum())}/{FULL_NODE_MULT.numel()})")
    assert FULL_NODE_SRC_IDX.min() >= 0 and FULL_NODE_SRC_IDX.max() < 94
    # Every one of the 94 input columns must be reachable by at least one (joint, channel) slot,
    # or "full" mode would be silently dropping input data instead of just reshaping it.
    reachable = set(FULL_NODE_SRC_IDX[FULL_NODE_MULT != 0].tolist())
    assert reachable == set(range(94)), f"full-feature layout does not cover all 94 columns: missing {set(range(94)) - reachable}"

    for node_features in ("coords", "full"):
        model = SkatingGCNClassifier(num_classes, node_features=node_features).to(DEVICE)
        logits = model(torch.from_numpy(X[:4]).to(DEVICE))
        assert logits.shape == (4, num_classes), logits.shape
        loss = nn.CrossEntropyLoss()(logits, torch.from_numpy(y[:4]).to(DEVICE))
        loss.backward()
        n_params = sum(p.numel() for p in model.parameters())
        print(f"[GCN] node_features={node_features}: forward OK: logits {tuple(logits.shape)} | "
              f"loss {loss.item():.3f} | backward OK | params {n_params/1e6:.2f}M")
    print(f"[GCN] num_classes={num_classes} matches softmax units; argmax range in [0,{num_classes-1}]")
    print("[GCN] SMOKE PASSED")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path)
    parser.add_argument("--coarse", action="store_true",
                        help="collapse jump rotations into action-level classes (11 instead of 28)")
    parser.add_argument("--layers", type=int, default=N_ATTN_BLOCKS,
                        help=f"number of Transformer encoder layers (default {N_ATTN_BLOCKS})")
    parser.add_argument("--node-features", choices=["coords", "full"], default="coords",
                        help="'coords': (x,y) only (default). 'full': all 94 dims scattered "
                             "onto their anatomically natural node slot -- see module docstring")
    parser.add_argument("--gcn-layers", type=int, default=len(GCN_CHANNELS),
                        help=f"number of spatial GCN blocks (default {len(GCN_CHANNELS)}); "
                             "channel widths ramp 64->128 across however many layers you pick")
    parser.add_argument("--noise-std", type=float, default=0.0,
                        help="train-only Gaussian noise std added to the standardized input "
                             "(0.0 = off). Val/test are never noised.")
    parser.add_argument("--smoke", action="store_true", help="self-test on synthetic data")
    args = parser.parse_args()
    if args.smoke:
        smoke()
        return 0
    if args.data_dir is None or not (args.data_dir / "train_features.pkl").exists():
        raise SystemExit(f"No processed data at {args.data_dir}. Run the pipeline first.")
    train(args.data_dir, coarse=args.coarse, num_layers=args.layers,
          node_features=args.node_features, gcn_layers=args.gcn_layers, noise_std=args.noise_std)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())