source/train.py

from __future__ import annotations

import json
import random
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import TemporalAutoencoder, parameter_count
from safetensors.torch import save_file
from sklearn.metrics import (
    average_precision_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

PROJECT_DIR = Path(__file__).resolve().parent
SOURCE_DATA_DIR = PROJECT_DIR.parent / "edge-sentinel-ml" / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "temporal-autoencoder"
CHANNELS = [
    "temperature",
    "pressure",
    "vibration",
    "current",
    "flow",
    "packet_rate",
    "command_rate",
    "actuator_position",
    "flow_actuator_residual",
    "power_proxy",
]
WINDOW_SIZE = 32
STRIDE = 8


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def load_frame(name: str) -> pd.DataFrame:
    path = SOURCE_DATA_DIR / f"{name}.parquet"
    if not path.exists():
        raise FileNotFoundError(
            f"{path} is missing. Run edge-sentinel-ml/generate_data.py first."
        )
    return pd.read_parquet(path)


def windows_from_frame(
    frame: pd.DataFrame,
    mean: np.ndarray,
    scale: np.ndarray,
) -> tuple[torch.Tensor, torch.Tensor]:
    windows = []
    labels = []
    for _, device in frame.groupby("device_id", sort=False):
        values = (device[CHANNELS].to_numpy(dtype=np.float32) - mean) / scale
        anomaly = device["label"].to_numpy(dtype=np.int64)
        for start in range(0, len(device) - WINDOW_SIZE + 1, STRIDE):
            stop = start + WINDOW_SIZE
            windows.append(values[start:stop].T)
            labels.append(int(anomaly[start:stop].any()))
    return (
        torch.from_numpy(np.stack(windows).astype(np.float32)),
        torch.tensor(labels, dtype=torch.long),
    )


@torch.inference_mode()
def reconstruction_scores(
    model: TemporalAutoencoder,
    windows: torch.Tensor,
) -> np.ndarray:
    model.eval()
    loader = DataLoader(TensorDataset(windows), batch_size=512, shuffle=False)
    scores = []
    for (batch,) in loader:
        reconstruction = model(batch)
        scores.extend(
            F.mse_loss(reconstruction, batch, reduction="none").mean((1, 2)).tolist()
        )
    return np.asarray(scores)


def best_threshold(labels: np.ndarray, scores: np.ndarray) -> tuple[float, float]:
    candidates = np.quantile(scores, np.linspace(0.5, 0.999, 500))
    ranked = [
        (f1_score(labels, scores >= threshold, zero_division=0), threshold)
        for threshold in candidates
    ]
    f1, threshold = max(ranked)
    return float(threshold), float(f1)


def metrics(labels: np.ndarray, scores: np.ndarray, threshold: float) -> dict:
    predictions = scores >= threshold
    matrix = confusion_matrix(labels, predictions, labels=[0, 1])
    return {
        "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)),
        "false_positive_rate": float(matrix[0, 1] / max(1, matrix[0].sum())),
        "confusion_matrix": matrix.tolist(),
    }


def main() -> None:
    seed_everything(2026)
    train_frame = load_frame("train")
    validation_frame = load_frame("validation")
    test_frame = load_frame("test")
    normal_train = train_frame[train_frame["label"] == 0]
    mean = normal_train[CHANNELS].to_numpy(dtype=np.float32).mean(axis=0)
    scale = normal_train[CHANNELS].to_numpy(dtype=np.float32).std(axis=0)
    scale = np.maximum(scale, 1e-4)
    train_windows, train_labels = windows_from_frame(train_frame, mean, scale)
    validation_windows, validation_labels = windows_from_frame(
        validation_frame,
        mean,
        scale,
    )
    test_windows, test_labels = windows_from_frame(test_frame, mean, scale)
    normal_windows = train_windows[train_labels == 0]
    loader = DataLoader(
        TensorDataset(normal_windows),
        batch_size=128,
        shuffle=True,
        generator=torch.Generator().manual_seed(2026),
    )
    model = TemporalAutoencoder(channels=len(CHANNELS))
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.002, weight_decay=0.001)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=35)
    best_validation_auc = -1.0
    best_state = None
    trackio.init(
        project="edge-sentinel-neural",
        name="temporal-autoencoder-v1",
        config={
            "parameters": parameter_count(model),
            "window_size": WINDOW_SIZE,
            "stride": STRIDE,
            "channels": len(CHANNELS),
            "normal_train_windows": len(normal_windows),
        },
    )
    for epoch in range(1, 36):
        model.train()
        running_loss = 0.0
        examples = 0
        for (batch,) in loader:
            noisy = batch + torch.randn_like(batch) * 0.025
            reconstruction = model(noisy)
            loss = F.mse_loss(reconstruction, batch)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * len(batch)
            examples += len(batch)
        scheduler.step()
        validation_scores = reconstruction_scores(model, validation_windows)
        validation_auc = roc_auc_score(validation_labels.numpy(), validation_scores)
        if validation_auc > best_validation_auc:
            best_validation_auc = validation_auc
            best_state = {
                key: value.detach().cpu().clone()
                for key, value in model.state_dict().items()
            }
        trackio.log(
            {
                "epoch": epoch,
                "train_reconstruction_mse": running_loss / examples,
                "validation_roc_auc": validation_auc,
                "learning_rate": scheduler.get_last_lr()[0],
            }
        )
    trackio.finish()
    if best_state is None:
        sys.exit("Training did not produce a checkpoint.")
    model.load_state_dict(best_state)
    validation_scores = reconstruction_scores(model, validation_windows)
    test_scores = reconstruction_scores(model, test_windows)
    threshold, validation_f1 = best_threshold(
        validation_labels.numpy(),
        validation_scores,
    )
    results = {
        "model": "Edge Sentinel Temporal Autoencoder",
        "parameters": parameter_count(model),
        "channels": CHANNELS,
        "window_size": WINDOW_SIZE,
        "stride": STRIDE,
        "normal_train_windows": len(normal_windows),
        "validation_windows": len(validation_windows),
        "test_windows": len(test_windows),
        "best_validation_roc_auc": float(best_validation_auc),
        "validation_threshold_f1": validation_f1,
        "threshold": threshold,
        "validation": metrics(
            validation_labels.numpy(),
            validation_scores,
            threshold,
        ),
        "test": metrics(test_labels.numpy(), test_scores, threshold),
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
    np.savez(
        ARTIFACT_DIR / "normalization.npz",
        mean=mean,
        scale=scale,
    )
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()