demo.py

"""ProbMon demo: full probabilistic condition monitoring pipeline.

Generates synthetic sensor data, runs Bayesian Online Change-Point Detection,
estimates Remaining Useful Life via Bayesian degradation modeling,
and produces visualizations and a summary report.
"""

import argparse
import os
import sys

import numpy as np

from probmon.sensors import TemperatureSensor, VibrationSensor, CurrentSensor
from probmon.changepoint import BOCDDetector
from probmon.degradation import BayesianDegradation
from probmon.plots import (
    plot_changepoint_detection,
    plot_rul_prediction,
    plot_multi_sensor_dashboard,
    plot_posterior_evolution,
)
from probmon.dashboard import generate_report


# Sensor configurations: (class, BOCD params, degradation params)
SENSOR_CONFIG = {
    "temp": {
        "class": TemperatureSensor,
        "label": "LED Temperature",
        "bocd": {"hazard_rate": 1 / 100, "mu0": 0, "kappa0": 1.0, "alpha0": 0.5, "beta0": 1.0, "threshold": 0.5},
        "failure_threshold": 95.0,
        "deg_prior": {"prior_rate_mean": 0.01, "prior_rate_std": 0.1},
    },
    "vibration": {
        "class": VibrationSensor,
        "label": "Bearing Vibration",
        "bocd": {"hazard_rate": 1 / 80, "mu0": 0, "kappa0": 1.0, "alpha0": 0.5, "beta0": 0.5, "threshold": 0.5},
        "failure_threshold": 5.0,
        "deg_prior": {"prior_rate_mean": 0.002, "prior_rate_std": 0.01},
    },
    "current": {
        "class": CurrentSensor,
        "label": "Leakage Current",
        "bocd": {"hazard_rate": 1 / 80, "mu0": 0, "kappa0": 1.0, "alpha0": 0.5, "beta0": 0.5, "threshold": 0.5},
        "failure_threshold": 8.0,
        "deg_prior": {"prior_rate_mean": 0.005, "prior_rate_std": 0.05},
    },
}


def run_sensor(
    sensor_key: str,
    n_points: int,
    output_dir: str,
    seed: int = 42,
) -> dict:
    """Run the full pipeline for a single sensor."""
    cfg = SENSOR_CONFIG[sensor_key]
    label = cfg["label"]

    print(f"\n--- {label} ---")

    # 1. Generate data
    sensor = cfg["class"]()
    t, values, gt_cps = sensor.generate(n_points=n_points, seed=seed)
    print(f"  Generated {n_points} points, {len(gt_cps)} ground truth change-points")

    # 2. BOCD on differenced data (stationarizes trends)
    diff_values = np.diff(values, prepend=values[0])
    detector = BOCDDetector(**cfg["bocd"])
    rl_matrix, cp_signal, detected_cps = detector.detect(diff_values)
    scores = detector.score(detected_cps, gt_cps, tolerance=20)
    scores["n_detected"] = len(detected_cps)
    scores["n_ground_truth"] = len(gt_cps)
    print(f"  Detected {len(detected_cps)} change-points | F1={scores['f1']:.2f}")

    # 3. Degradation model
    deg_model = BayesianDegradation(
        failure_threshold=cfg["failure_threshold"],
        **cfg["deg_prior"],
    )
    deg_model.fit(t, values)
    rul_mean, rul_std, rul_samples = deg_model.predict_rul()
    future_t, mean_traj, std_traj = deg_model.predict_trajectory(future_steps=200)
    post_means, post_stds = deg_model.get_posterior_history()
    print(f"  RUL estimate: {rul_mean:.1f} +/- {rul_std:.1f} time steps")

    # 4. Plots
    plot_changepoint_detection(
        t, values, rl_matrix, cp_signal, detected_cps, gt_cps,
        sensor_name=label,
        output_path=os.path.join(output_dir, f"bocd_{sensor_key}.png"),
    )
    plot_rul_prediction(
        t, values, future_t, mean_traj, std_traj,
        failure_threshold=cfg["failure_threshold"],
        rul_mean=rul_mean, rul_std=rul_std,
        sensor_name=label,
        output_path=os.path.join(output_dir, f"rul_{sensor_key}.png"),
    )
    plot_posterior_evolution(
        t, post_means, post_stds,
        sensor_name=label,
        output_path=os.path.join(output_dir, f"posterior_{sensor_key}.png"),
    )

    return {
        "timestamps": t,
        "values": values,
        "detected_cps": detected_cps,
        "ground_truth_cps": gt_cps,
        "cp_signal": cp_signal,
        "scores": scores,
        "rul_mean": rul_mean,
        "rul_std": rul_std,
        "failure_threshold": cfg["failure_threshold"],
    }


def main():
    parser = argparse.ArgumentParser(
        description="ProbMon: Probabilistic Condition Monitoring Demo"
    )
    parser.add_argument(
        "--sensor",
        choices=["temp", "vibration", "current", "all"],
        default="all",
        help="Sensor to analyze (default: all)",
    )
    parser.add_argument(
        "--points", type=int, default=500,
        help="Number of data points to generate (default: 500)",
    )
    parser.add_argument(
        "--output-dir", type=str, default="output",
        help="Directory for output plots (default: output)",
    )
    parser.add_argument(
        "--seed", type=int, default=42,
        help="Random seed (default: 42)",
    )
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    np.random.seed(args.seed)

    print("=" * 60)
    print("  ProbMon: Probabilistic Condition Monitoring")
    print("=" * 60)

    sensors = (
        list(SENSOR_CONFIG.keys()) if args.sensor == "all" else [args.sensor]
    )

    all_results = {}
    sensor_scores = {}
    rul_results = {}

    for key in sensors:
        res = run_sensor(key, args.points, args.output_dir, seed=args.seed)
        label = SENSOR_CONFIG[key]["label"]
        all_results[label] = res
        sensor_scores[label] = res["scores"]
        rul_results[label] = {
            "rul_mean": res["rul_mean"],
            "rul_std": res["rul_std"],
            "failure_threshold": res["failure_threshold"],
        }

    # Multi-sensor dashboard (if multiple sensors)
    if len(all_results) > 1:
        plot_multi_sensor_dashboard(
            all_results,
            output_path=os.path.join(args.output_dir, "dashboard.png"),
        )

    # Text report
    report = generate_report(sensor_scores, rul_results)
    print("\n")
    print(report)

    report_path = os.path.join(args.output_dir, "report.txt")
    with open(report_path, "w", encoding="utf-8") as f:
        f.write(report)

    print(f"\nPlots saved to: {os.path.abspath(args.output_dir)}/")


if __name__ == "__main__":
    main()