14_07_expt/processed/model_ctm.py

"""Continuous Thought Machine with Adaptive Computation Time for figure-skating classification.

Architecture: Conv1D backbone (jointly trained) + CTM thinking loop with ACT halting.
The model learns to halt early on easy samples (spins) and think longer on hard ones
(ambiguous jump rotations), instead of running a fixed number of iterations.

Usage:
    python model_ctm.py --data-dir /path/to/processed
    python model_ctm.py --data-dir /path/to/processed --coarse
"""

from __future__ import annotations

import argparse
import math
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
except ImportError:
    import labels as labels_mod

from ctm_components import NeuronModel, SynapseNet, compute_synchronization
from model import (
    ConvBlock,
    load_split,
    resolve_num_classes,
    assert_label_consistency,
    report_metrics,
)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

EPOCHS = 100
BATCH_SIZE = 64
LEARNING_RATE = 3e-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"

CTM_MAX_ITERATIONS = 30
CTM_D_MODEL = 256
CTM_D_BACKBONE = 384
CTM_MEMORY_LENGTH = 50
CTM_N_SYNCH = 64
CTM_NUM_HEADS = 8
CTM_DROPOUT = 0.1
ACT_EPSILON = 0.01
PONDER_LAMBDA = 0.01


# ---------------------------------------------------------------------------
# Conv backbone (jointly trained feature encoder)
# ---------------------------------------------------------------------------

class ConvBackbone(nn.Module):
    """Conv1D encoder: (B, T, in_features) -> (B, T, d_backbone)."""

    def __init__(self, in_features: int, d_backbone: int = 384):
        super().__init__()
        self.stem = nn.Conv1d(in_features, 128, 3, padding="same")
        self.stem_bn = nn.BatchNorm1d(128)
        self.cb1 = ConvBlock(128, 192, 3)
        self.cb2 = ConvBlock(192, 256, 3)
        self.cb3 = ConvBlock(256, d_backbone, 5)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.transpose(1, 2)
        x = F.relu(self.stem(x))
        x = self.stem_bn(x)
        x = self.cb3(self.cb2(self.cb1(x)))
        return x.transpose(1, 2)


# ---------------------------------------------------------------------------
# Skating CTM + ACT
# ---------------------------------------------------------------------------

class SkatingCTM(nn.Module):
    """Conv backbone + CTM thinking loop with Adaptive Computation Time.

    Instead of a fixed iteration count, each sample halts independently once
    the learned halt probability accumulates past 1 - epsilon. Easy samples
    (spins) halt early; ambiguous jumps get more thinking steps.
    """

    def __init__(
        self,
        in_features: int,
        num_classes: int,
        d_model: int = CTM_D_MODEL,
        d_backbone: int = CTM_D_BACKBONE,
        num_heads: int = CTM_NUM_HEADS,
        memory_length: int = CTM_MEMORY_LENGTH,
        max_iterations: int = CTM_MAX_ITERATIONS,
        n_synch: int = CTM_N_SYNCH,
        dropout: float = CTM_DROPOUT,
        act_epsilon: float = ACT_EPSILON,
        use_act: bool = True,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_classes = num_classes
        self.max_iterations = max_iterations
        self.n_synch = n_synch
        self.act_epsilon = act_epsilon
        self.use_act = use_act

        self.backbone = ConvBackbone(in_features, d_backbone)
        self.kv_proj = nn.Sequential(
            nn.Linear(d_backbone, d_model),
            nn.LayerNorm(d_model),
        )

        self.nlm = NeuronModel(memory_length, d_model, dropout=dropout)
        self.synapses = SynapseNet(d_model, d_model, dropout=dropout)
        self.attention = nn.MultiheadAttention(
            d_model, num_heads, dropout=dropout, batch_first=True,
        )

        synch_rep = (n_synch * (n_synch + 1)) // 2
        self.q_proj = nn.Linear(synch_rep, d_model)
        self.output_proj = nn.Linear(synch_rep, num_classes)
        self.decay_action = nn.Parameter(torch.zeros(synch_rep))
        self.decay_out = nn.Parameter(torch.zeros(synch_rep))

        self.start_act = nn.Parameter(torch.zeros(d_model).uniform_(-0.1, 0.1))
        self.start_trace = nn.Parameter(torch.zeros(d_model, memory_length).uniform_(-0.1, 0.1))

        # ACT halting unit — bias initialized negative so initial halt prob is
        # low (~0.05) and the model starts by using ~20 steps, learning to halt
        # earlier as training progresses.
        self.halt_proj = nn.Linear(d_model, 1)
        nn.init.constant_(self.halt_proj.bias, -3.0)

        self.register_buffer("idx_left_action", torch.arange(d_model - n_synch, d_model))
        self.register_buffer("idx_right_action", torch.arange(d_model - n_synch, d_model))
        self.register_buffer("idx_left_out", torch.arange(0, n_synch))
        self.register_buffer("idx_right_out", torch.arange(0, n_synch))

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Returns:
            preds:        (B, num_classes, T_actual) per-step predictions
            halt_weights: (B, T_actual) ACT weights per step (sum to ~1 per sample)
            ponder_cost:  scalar, mean steps taken across batch
            n_steps:      (B,) per-sample step count
        """
        features = self.backbone(x)
        B = features.size(0)
        kv = self.kv_proj(features)

        state_trace = self.start_trace.unsqueeze(0).expand(B, -1, -1).clone()
        activated_state = self.start_act.unsqueeze(0).expand(B, -1).clone()

        decay_action = torch.exp(-self.decay_action.clamp(0, 15)).unsqueeze(0).expand(B, -1)
        decay_out = torch.exp(-self.decay_out.clamp(0, 15)).unsqueeze(0).expand(B, -1)

        _, ema_n_out, ema_d_out = compute_synchronization(
            activated_state, None, None, decay_out,
            self.n_synch, self.idx_left_out, self.idx_right_out,
        )

        preds_list: list[torch.Tensor] = []
        halt_list: list[torch.Tensor] = []
        cumulative_p = torch.zeros(B, device=x.device)
        halted = torch.zeros(B, dtype=torch.bool, device=x.device)
        n_steps = torch.zeros(B, device=x.device)
        ema_n_act = ema_d_act = None

        for t in range(self.max_iterations):
            # --- thinking step ---
            sync_act, ema_n_act, ema_d_act = compute_synchronization(
                activated_state, ema_n_act, ema_d_act, decay_action,
                self.n_synch, self.idx_left_action, self.idx_right_action,
            )
            q = self.q_proj(sync_act).unsqueeze(1)
            attn_out, _ = self.attention(q, kv, kv, need_weights=False)
            attn_out = attn_out.squeeze(1)

            state = self.synapses(torch.cat([attn_out, activated_state], dim=-1))
            state_trace = torch.cat([state_trace[:, :, 1:], state.unsqueeze(-1)], dim=-1)
            activated_state = self.nlm(state_trace)

            sync_out, ema_n_out, ema_d_out = compute_synchronization(
                activated_state, ema_n_out, ema_d_out, decay_out,
                self.n_synch, self.idx_left_out, self.idx_right_out,
            )
            preds_list.append(self.output_proj(sync_out))

            if not self.use_act:
                # Fixed-iteration mode: no halting: every step contributes equally to the
                # loss/prediction and every sample always runs exactly max_iterations steps.
                halt_list.append(torch.full((B,), 1.0 / self.max_iterations, device=x.device))
                n_steps += 1.0
                continue

            # --- ACT halting ---
            p = torch.sigmoid(self.halt_proj(activated_state)).squeeze(-1)
            still_running = ~halted
            new_cumulative = cumulative_p + p

            halt_now = still_running & (new_cumulative >= 1.0 - self.act_epsilon)
            remainder = 1.0 - cumulative_p
            w = torch.where(halt_now, remainder, p) * still_running.float()
            halt_list.append(w)

            cumulative_p = torch.where(halt_now | halted, cumulative_p, new_cumulative)
            n_steps += still_running.float()
            halted = halted | halt_now

            if halted.all():
                break

        # Assign remainder to samples that never halted within max_iterations
        if self.use_act and not halted.all():
            still_running = ~halted
            halt_list[-1] = halt_list[-1] + (1.0 - cumulative_p) * still_running.float()

        preds = torch.stack(preds_list, dim=-1)
        halt_weights = torch.stack(halt_list, dim=-1)
        ponder_cost = n_steps.mean()

        return preds, halt_weights, ponder_cost, n_steps


# ---------------------------------------------------------------------------
# Loss
# ---------------------------------------------------------------------------

def ctm_act_loss(
    preds: torch.Tensor,
    targets: torch.Tensor,
    halt_weights: torch.Tensor,
    ponder_cost: torch.Tensor,
    class_weight: torch.Tensor | None = None,
    ponder_lambda: float = PONDER_LAMBDA,
) -> torch.Tensor:
    """Halt-weighted CE across thinking steps + ponder cost regularization."""
    _, _, T = preds.shape
    total = preds.new_zeros(())
    for t in range(T):
        ce = F.cross_entropy(preds[:, :, t], targets, weight=class_weight, reduction="none")
        total = total + (halt_weights[:, t] * ce).mean()
    return total + ponder_lambda * ponder_cost


# ---------------------------------------------------------------------------
# Train / evaluate
# ---------------------------------------------------------------------------

@torch.no_grad()
def predict(model: nn.Module, X: np.ndarray) -> tuple[np.ndarray, float]:
    """Returns (predicted_classes, mean_steps)."""
    model.eval()
    out = []
    total_steps = 0.0
    total_n = 0
    for i in range(0, len(X), 256):
        xb = torch.from_numpy(X[i : i + 256]).to(DEVICE)
        preds, halt_weights, _, n_steps = model(xb)
        weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1)
        out.append(weighted.argmax(1).cpu().numpy())
        total_steps += n_steps.sum().item()
        total_n += xb.size(0)
    classes = np.concatenate(out) if out else np.array([], dtype=np.int64)
    return classes, total_steps / max(total_n, 1)


def train(
    data_dir: Path,
    coarse: bool = False,
    max_iterations: int = CTM_MAX_ITERATIONS,
    memory_length: int = CTM_MEMORY_LENGTH,
    use_act: bool = True,
    ponder_lambda: float = PONDER_LAMBDA,
    tag: str = "CTM-ACT",
) -> 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]
    print(f"[{tag}] label space: {'COARSE' 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"[{tag}] device={DEVICE} | in_features={in_features}")
    print(f"[{tag}] max_iterations={max_iterations} | memory={memory_length} | use_act={use_act} | ponder_lambda={ponder_lambda}")
    print(f"[{tag}] shapes: train={Xtr.shape} val={Xva.shape} test={Xte.shape}")

    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)
    class_weight = weight.to(DEVICE)

    model = SkatingCTM(
        in_features, num_classes,
        max_iterations=max_iterations, memory_length=memory_length, use_act=use_act,
    ).to(DEVICE)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[{tag}] params: {n_params / 1e6:.2f}M")
    optimizer = torch.optim.AdamW(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):
        model.train()
        tr_loss, tr_correct, tr_n = 0.0, 0, 0
        epoch_steps = 0.0
        for xb, yb in tr_loader:
            xb, yb = xb.to(DEVICE), yb.to(DEVICE)
            preds, halt_weights, ponder_cost, n_steps = model(xb)
            loss = ctm_act_loss(preds, yb, halt_weights, ponder_cost, class_weight, ponder_lambda=ponder_lambda)
            if not torch.isfinite(loss):
                continue
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
            optimizer.step()
            tr_loss += loss.item() * xb.size(0)
            weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1)
            tr_correct += (weighted.argmax(1) == yb).sum().item()
            tr_n += xb.size(0)
            epoch_steps += n_steps.sum().item()

        model.eval()
        va_loss, va_correct, va_n = 0.0, 0, 0
        va_steps = 0.0
        with torch.no_grad():
            for xb, yb in va_loader:
                xb, yb = xb.to(DEVICE), yb.to(DEVICE)
                preds, halt_weights, ponder_cost, n_steps = model(xb)
                loss = ctm_act_loss(preds, yb, halt_weights, ponder_cost, class_weight, ponder_lambda=ponder_lambda)
                va_loss += loss.item() * xb.size(0)
                weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1)
                va_correct += (weighted.argmax(1) == yb).sum().item()
                va_n += xb.size(0)
                va_steps += n_steps.sum().item()

        tr_loss /= tr_n
        tr_acc = tr_correct / tr_n
        va_loss /= va_n
        va_acc = va_correct / va_n
        mean_tr_steps = epoch_steps / tr_n
        mean_va_steps = va_steps / va_n

        va_preds, _ = predict(model, Xva)
        va_f1 = f1_score(
            yva, va_preds,
            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"[{tag}] epoch {epoch:3d} | train loss {tr_loss:.3f} acc {tr_acc:.3f} "
                f"| val loss {va_loss:.3f} acc {va_acc:.3f} f1 {va_f1:.3f} "
                f"| steps tr={mean_tr_steps:.1f} va={mean_va_steps:.1f}"
            )

        if since_improved >= EARLY_STOP_PATIENCE:
            print(f"[{tag}] 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[{tag}] restored best model (val macro-F1 = {best_val_f1:.3f})")

    te_preds, te_mean_steps = predict(model, Xte)
    print(f"[{tag}] test mean thinking steps: {te_mean_steps:.1f}")
    metrics = report_metrics(yte, te_preds, taxonomy)
    metrics["mean_test_steps"] = te_mean_steps
    metrics["n_params"] = n_params
    suffix = f"_{tag.lower()}" if tag != "CTM-ACT" else ""
    ckpt_name = f"model_ctm{suffix}{'_coarse' if coarse else ''}.pt"
    torch.save(
        {
            "state_dict": model.state_dict(),
            "num_classes": num_classes,
            "in_features": in_features,
            "taxonomy": taxonomy,
            "coarse": coarse,
            "feature_mean": mu,
            "feature_std": sd,
            "max_iterations": max_iterations,
            "memory_length": memory_length,
            "use_act": use_act,
            "act_epsilon": ACT_EPSILON,
        },
        data_dir / ckpt_name,
    )
    print(f"\n[{tag}] saved -> {data_dir / ckpt_name}")
    return metrics


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, required=True)
    parser.add_argument("--coarse", action="store_true",
                        help="collapse jump rotations into action-level classes (11 instead of 28)")
    parser.add_argument("--iterations", type=int, default=CTM_MAX_ITERATIONS,
                        help="thinking steps: max steps for ACT, exact fixed steps otherwise")
    parser.add_argument("--no-act", action="store_true",
                        help="disable adaptive halting: run exactly --iterations steps, uniform loss/prediction averaging")
    parser.add_argument("--memory-length", type=int, default=CTM_MEMORY_LENGTH)
    parser.add_argument("--ponder-lambda", type=float, default=PONDER_LAMBDA)
    parser.add_argument("--tag", type=str, default=None, help="log/checkpoint tag, e.g. CTM-10")
    args = parser.parse_args()
    if not (args.data_dir / "train_features.pkl").exists():
        raise SystemExit(f"No processed data at {args.data_dir}. Run the pipeline first.")
    use_act = not args.no_act
    tag = args.tag or (f"CTM-{args.iterations}" if not use_act else "CTM-ACT")
    train(
        args.data_dir, coarse=args.coarse,
        max_iterations=args.iterations, memory_length=args.memory_length,
        use_act=use_act, ponder_lambda=(0.0 if not use_act else args.ponder_lambda),
        tag=tag,
    )
    return 0


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