test.py

#!/usr/bin/env python
"""
Evaluate a saved GNSS threat detector on the cross-test set.

Loads a model saved by train.py, applies the saved preprocessors,
and reports accuracy and per-class metrics.

The model directory must contain:
  {model}.keras or {model}.joblib  — trained model
  imputer.joblib, scaler.joblib    — preprocessors (neural-net models only)
  metrics.json                     — model name and window_size written by train.py

Usage
-----
python test.py --model-dir ./runs/bilstm_w10
python test.py --model-dir ./runs/rf_w20 --output ./runs/rf_w20/cross_test.json
"""

import argparse
import json
import time
from pathlib import Path
from typing import Dict, Optional

import joblib
import numpy as np
import pandas as pd

from const import (
    FEATURES,
    NEURAL_NETS,
    TEST_CSV,
)

from utils import create_training_data, compute_metrics, print_metrics

def run(
    model_dir: Path,
    output: Optional[Path] = None,
) -> Dict:
    """Evaluate the model in *model_dir* on the cross-test set.

    Parameters
    ----------
    model_dir:
        Directory produced by ``train.py --output-dir``.
    output:
        If set, write metrics to this JSON file.

    Returns
    -------
    dict with accuracy, per-class precision/recall/F1, and confusion matrix.
    """
    model_dir = Path(model_dir)

    # Read model name and window_size saved by train.py
    config     = json.loads((model_dir / "metrics.json").read_text())
    model_name = config["model"]
    window_size = config["window_size"]

    if not TEST_CSV.exists():
        raise FileNotFoundError(f"{TEST_CSV} not found — run preprocess.py first")

    print("Reading cross-test data …")
    test_df = pd.read_csv(TEST_CSV, parse_dates=["ts"])
    X_test, y_test = create_training_data(test_df, window_size=window_size, features=FEATURES)

    # Apply saved preprocessors (neural nets only)
    is_nn = model_name in NEURAL_NETS
    if is_nn:
        imputer = joblib.load(model_dir / "imputer.joblib")
        scaler  = joblib.load(model_dir / "scaler.joblib")
        X_test  = np.nan_to_num(scaler.transform(imputer.transform(X_test)))

    # Load model and predict
    print(f"Evaluating {model_name.upper()} (window={window_size}s) …")
    model  = _load_model(model_dir, model_name)

    t0 = time.perf_counter()
    y_pred = _predict(model, model_name, X_test, window_size)
    avg_inference_ms = (time.perf_counter() - t0) / len(X_test) * 1000

    cross_metrics = {"cross_" + k: v for k, v in compute_metrics(y_test, y_pred).items()}
    cross_metrics["cross_avg_inference_ms"] = avg_inference_ms

    # Merge cross-test results into metrics.json (keeps existing val_* keys)
    metrics_path = model_dir / "metrics.json"
    metrics = json.loads(metrics_path.read_text())
    metrics.update(cross_metrics)
    metrics_path.write_text(json.dumps(metrics, indent=2))

    if output is not None:
        Path(output).write_text(json.dumps(metrics, indent=2))

    return metrics

# ---------------------------------------------------------------------------
# Model loading and prediction
# ---------------------------------------------------------------------------

def _load_model(model_dir: Path, model_name: str):
    """Load a model saved by train.py."""
    if model_name in NEURAL_NETS:
        import tensorflow as tf
        return tf.keras.models.load_model(model_dir / f"{model_name}.keras")
    return joblib.load(model_dir / f"{model_name}.joblib")


def _predict(model, model_name: str, X: np.ndarray, window_size: int) -> np.ndarray:
    """Run inference, reshaping X for neural-net models."""
    if model_name in NEURAL_NETS:
        n_features = X.shape[1] // window_size
        X_seq = X.reshape(-1, window_size, n_features)
        return np.argmax(model.predict(X_seq, verbose=0), axis=1)
    return model.predict(X)

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main(argv=None) -> None:
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument(
        "--model-dir", required=True, metavar="DIR",
        help="Directory containing the saved model (output of train.py)",
    )
    p.add_argument(
        "--output", default=None, metavar="FILE",
        help="Save metrics to this JSON file (optional)",
    )
    args = p.parse_args(argv)

    metrics = run(
        model_dir=Path(args.model_dir),
        output=Path(args.output) if args.output else None,
    )

    print(f"\n  Model : {metrics['model'].upper()}  (window={metrics['window_size']}s)")
    print_metrics(metrics, prefix="cross_", title="Cross-test results")
    print(f"\n  Avg. inference time : {metrics['cross_avg_inference_ms']:.4f} ms/sample")


if __name__ == "__main__":
    main()