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()