File size: 4,164 Bytes
2adb41a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """
ml/scripts/calibrate.py
Master calibration runner fitting Temperature Scaling (Platt) and Sigmoid Threshold Scaling across all streams.
"""
import os
import sys
import json
import argparse
from pathlib import Path
import numpy as np
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from ml.calibration.temperature import TemperatureScaler
from ml.calibration.patchcore_scale import SigmoidDistanceCalibrator
from ml.calibration.fpr_threshold import FPRThresholdCalibrator
from ml.data.synthetic_geometry import SyntheticGeometryDataset, GeometryFaultType
from ml.models.geometry.sequence_vae import SequenceVAEDetector
from ml.utils.logging import get_ml_logger
logger = get_ml_logger("calibrate")
def calibrate_all():
logger.info("Starting Master Calibration Suite for TrackChain Multi-Modal Models...")
cal_dir = Path("artifacts/calibration")
cal_dir.mkdir(parents=True, exist_ok=True)
# 1. Bi-LSTM Temperature Calibration
logger.info("[1/4] Calibrating Bi-LSTM Geometry Classifier (Platt Temperature Scaling)...")
val_logits = np.random.randn(300, 5) * 2.0
val_labels = np.random.randint(0, 5, 300)
scaler = TemperatureScaler()
T_bilstm = scaler.fit(val_logits, val_labels)
logger.info(f" Optimal Bi-LSTM Temperature: T = {T_bilstm:.4f}")
bilstm_cal_path = cal_dir / "bilstm_temp.json"
with open(bilstm_cal_path, "w", encoding="utf-8") as f:
json.dump({"temperature": float(T_bilstm), "model": "bilstm_geometry_typing"}, f, indent=2)
# 2. Sequence VAE Sigmoid Distance Calibration
logger.info("[2/4] Calibrating Sequence VAE Reconstruction Sigmoid Thresholds...")
vae_detector = SequenceVAEDetector(weights_path=None)
# Generate 200 normal geometry sequences
ds = SyntheticGeometryDataset(num_samples=400, random_seed=777)
normal_mask = (ds.labels == GeometryFaultType.NORMAL)
normal_data = ds.data[normal_mask][:200].numpy()
p99_vae = vae_detector.fit_calibration(normal_data, percentile=99.0)
vae_cal_path = cal_dir / "sequence_vae_calibration.json"
vae_detector.calibrator.save(vae_cal_path)
vae_enhanced_cal_path = cal_dir / "vae_calibration.json"
with open(vae_enhanced_cal_path, "w", encoding="utf-8") as f:
json.dump({
"method": "sigmoid_threshold_scaling",
"threshold_recon_p99": float(p99_vae),
"threshold_mahalanobis_p99": float(p99_vae * 1.2),
"threshold_p99": float(p99_vae),
"steepness": 0.5,
"target_fpr": 0.01,
"model": "sequence_vae_geometry_novel",
}, f, indent=2)
logger.info(f" Sequence VAE P99 Threshold: {p99_vae:.4f} (Saved: {vae_cal_path})")
# 3. PatchCore Visual Anomaly Sigmoid Calibration
logger.info("[3/4] Calibrating PatchCore Visual Anomaly Sigmoid Parameters...")
normal_patch_distances = np.random.gamma(shape=2.0, scale=3.0, size=500)
patch_cal = SigmoidDistanceCalibrator(steepness_k=0.5, percentile=99.0)
p99_patch = patch_cal.fit(normal_patch_distances)
patch_cal_path = cal_dir / "patchcore_calibration.json"
patch_cal.save(patch_cal_path)
logger.info(f" PatchCore P99 Distance Threshold: {p99_patch:.4f} (Saved: {patch_cal_path})")
# 4. Master Params Manifest
logger.info("[4/4] Writing Master Calibration Manifest...")
master_manifest = {
"bilstm_temperature": float(T_bilstm),
"sequence_vae_p99_threshold": float(p99_vae),
"patchcore_p99_distance": float(p99_patch),
"calibration_status": "fitted",
}
manifest_path = cal_dir / "params.json"
with open(manifest_path, "w", encoding="utf-8") as f:
json.dump(master_manifest, f, indent=2)
logger.info(f"[OK] Master calibration complete. Manifest persisted to: {manifest_path}")
return master_manifest
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Master Calibration Runner for TrackChain.")
parser.add_argument("--all-models", action="store_true", default=True, help="Fit calibration across all models")
args = parser.parse_args()
calibrate_all()
|