source/train.py

from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import trackio
from sklearn.ensemble import HistGradientBoostingClassifier, IsolationForest
from sklearn.metrics import (
    average_precision_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from sklearn.preprocessing import RobustScaler

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "edge-sentinel-ensemble"
EXCLUDED = {"device_id", "time", "label", "anomaly_type", "split"}


def load_split(name: str) -> pd.DataFrame:
    return pd.read_parquet(DATA_DIR / f"{name}.parquet")


def metrics(labels: np.ndarray, scores: np.ndarray, threshold: float) -> dict:
    predictions = (scores >= threshold).astype(np.int8)
    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 best_threshold(labels: np.ndarray, scores: np.ndarray) -> tuple[float, float]:
    candidates = np.quantile(scores, np.linspace(0.7, 0.999, 300))
    best = (0.5, -1.0)
    for threshold in candidates:
        value = f1_score(labels, scores >= threshold, zero_division=0)
        if value > best[1]:
            best = (float(threshold), float(value))
    return best


def normalize(values: np.ndarray, reference: np.ndarray) -> np.ndarray:
    low, high = np.quantile(reference, [0.01, 0.99])
    return np.clip((values - low) / max(1e-9, high - low), 0, 1)


def main() -> None:
    train = load_split("train")
    validation = load_split("validation")
    test = load_split("test")
    features = [column for column in train.columns if column not in EXCLUDED]
    scaler = RobustScaler()
    x_train = scaler.fit_transform(train[features])
    x_validation = scaler.transform(validation[features])
    x_test = scaler.transform(test[features])
    y_train = train["label"].to_numpy()
    y_validation = validation["label"].to_numpy()
    y_test = test["label"].to_numpy()

    isolation = IsolationForest(
        n_estimators=300,
        max_samples=0.7,
        contamination="auto",
        random_state=42,
        n_jobs=-1,
    )
    isolation.fit(x_train[y_train == 0])
    supervised = HistGradientBoostingClassifier(
        learning_rate=0.08,
        max_iter=250,
        max_leaf_nodes=31,
        l2_regularization=0.2,
        class_weight="balanced",
        random_state=42,
    )
    supervised.fit(x_train, y_train)

    isolation_validation_raw = -isolation.score_samples(x_validation)
    isolation_test_raw = -isolation.score_samples(x_test)
    isolation_validation = normalize(isolation_validation_raw, isolation_validation_raw)
    isolation_test = normalize(isolation_test_raw, isolation_validation_raw)
    supervised_validation = supervised.predict_proba(x_validation)[:, 1]
    supervised_test = supervised.predict_proba(x_test)[:, 1]

    trackio.init(
        project="edge-sentinel-ml",
        name="device-held-out-ensemble-v1",
        config={
            "train_devices": 8,
            "validation_devices": 2,
            "test_devices": 2,
            "features": len(features),
        },
    )
    candidates = []
    for supervised_weight in np.linspace(0, 1, 21):
        validation_scores = (
            supervised_weight * supervised_validation
            + (1 - supervised_weight) * isolation_validation
        )
        threshold, validation_f1 = best_threshold(y_validation, validation_scores)
        candidates.append((validation_f1, supervised_weight, threshold))
        trackio.log(
            {
                "supervised_weight": supervised_weight,
                "validation_f1": validation_f1,
                "threshold": threshold,
            }
        )
    _, weight, threshold = max(candidates)
    validation_scores = weight * supervised_validation + (1 - weight) * isolation_validation
    test_scores = weight * supervised_test + (1 - weight) * isolation_test
    results = {
        "features": features,
        "ensemble_supervised_weight": float(weight),
        "threshold": threshold,
        "validation": metrics(y_validation, validation_scores, threshold),
        "test": metrics(y_test, test_scores, threshold),
        "isolation_test": metrics(
            y_test,
            isolation_test,
            best_threshold(y_validation, isolation_validation)[0],
        ),
        "supervised_test": metrics(
            y_test,
            supervised_test,
            best_threshold(y_validation, supervised_validation)[0],
        ),
    }
    trackio.log(
        {
            "test_roc_auc": results["test"]["roc_auc"],
            "test_average_precision": results["test"]["average_precision"],
            "test_f1": results["test"]["f1"],
            "test_recall": results["test"]["recall"],
            "test_false_positive_rate": results["test"]["false_positive_rate"],
        }
    )
    trackio.finish()

    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump(
        {
            "scaler": scaler,
            "isolation_forest": isolation,
            "supervised_model": supervised,
            "features": features,
            "weight": weight,
            "threshold": threshold,
        },
        ARTIFACT_DIR / "edge_sentinel.joblib",
        compress=3,
    )
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()