hackathon-ENGIN/ml_polish_benchmark.py

"""Repeated grouped validation for the final polished ENGIN architecture.

This is a compact reproducibility check for the exact models used by
``final_pipeline.py``. It evaluates complete held-out engines for five CV seeds
on clean validation data and the same data with exactly 5% spectral cells
masked. Training data remains clean.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score, f1_score

from benchmark_grouped import (
    FAULT_LABELS,
    LABELS,
    NOT_APPLICABLE,
    macro_f1,
    make_splits,
    ml_points,
    raw_score,
    validate_data,
)
from final_pipeline import (
    SEVERITY_CANDIDATE_ID,
    apply_ood_override,
    make_final_label_pipeline,
)
from robustness_grouped import mask_spectral_cells, missing_scenario_name
from severity_benchmark import (
    CANDIDATE_BY_ID,
    fit_candidate,
    predict_candidate,
    prepare_labeled_frame,
)

DEFAULT_SEEDS = [7, 21, 42, 77, 123]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", type=Path, default=Path("val.csv"))
    parser.add_argument("--output-dir", type=Path, default=Path("ml_polish_outputs"))
    parser.add_argument("--seeds", nargs="+", type=int, default=DEFAULT_SEEDS)
    parser.add_argument("--n-splits", type=int, default=5)
    parser.add_argument("--missing-rate", type=float, default=0.05)
    parser.add_argument("--mask-random-state", type=int, default=2026)
    parser.add_argument("--model-random-state", type=int, default=42)
    parser.add_argument("--n-jobs", type=int, default=-1)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    clean = pd.read_csv(args.data).reset_index(drop=True)
    validate_data(clean, args.n_splits)
    masked, mask_manifest = mask_spectral_cells(
        clean, args.missing_rate, args.mask_random_state
    )
    masked_name = missing_scenario_name(args.missing_rate)
    scenarios = {"clean": clean, masked_name: masked}
    y_label = clean["label"].reset_index(drop=True)
    y_severity = clean["severity"].reset_index(drop=True)
    true_fault = y_label.isin(FAULT_LABELS).to_numpy()
    candidate = CANDIDATE_BY_ID[SEVERITY_CANDIDATE_ID]
    run_rows: list[dict[str, object]] = []

    args.output_dir.mkdir(parents=True, exist_ok=True)
    mask_manifest.to_csv(args.output_dir / "mask_manifest.csv", index=False)

    for cv_seed in args.seeds:
        print(f"Running polished CV seed {cv_seed} ...", flush=True)
        splits = make_splits(
            clean, "stratified-group", args.n_splits, cv_seed
        )
        oof_label = {
            scenario: np.full(len(clean), "", dtype=object)
            for scenario in scenarios
        }
        oof_severity = {
            scenario: np.full(len(clean), NOT_APPLICABLE, dtype=object)
            for scenario in scenarios
        }
        ood_count = {scenario: 0 for scenario in scenarios}

        for fold, (train_idx, valid_idx) in enumerate(splits, start=1):
            train = clean.iloc[train_idx]
            train_labels = y_label.iloc[train_idx]
            label_pipeline = make_final_label_pipeline(
                args.model_random_state
            ).fit(train, train_labels)

            train_fault = train_labels.isin(FAULT_LABELS).to_numpy()
            train_with_labels = prepare_labeled_frame(
                train, train_labels.to_numpy(dtype=object)
            )
            transformer, estimator = fit_candidate(
                candidate=candidate,
                X_train_full=train_with_labels,
                train_fault_mask=train_fault,
                y_severity_fault=y_severity.iloc[train_idx].to_numpy(dtype=object)[
                    train_fault
                ],
                y_fault_label=train_labels.to_numpy(dtype=object)[train_fault],
                random_state=args.model_random_state + 1_002 + 20 * fold,
                n_jobs=args.n_jobs,
            )

            for scenario, frame in scenarios.items():
                valid = frame.iloc[valid_idx]
                base_label = label_pipeline.predict(valid).astype(object)
                label_features = label_pipeline.named_steps["features"].transform(
                    valid
                )
                predicted_label, override, _, _ = apply_ood_override(
                    base_label,
                    label_features,
                    valid["engine_id"],
                )
                oof_label[scenario][valid_idx] = predicted_label
                ood_count[scenario] += int(override.sum())

                valid_with_labels = prepare_labeled_frame(
                    valid, predicted_label
                )
                severity_all = predict_candidate(
                    candidate,
                    transformer,
                    estimator,
                    valid_with_labels,
                    predicted_label,
                )
                emitted = np.full(len(valid), NOT_APPLICABLE, dtype=object)
                predicted_fault = np.isin(predicted_label, FAULT_LABELS)
                emitted[predicted_fault] = severity_all[predicted_fault]
                oof_severity[scenario][valid_idx] = emitted

        for scenario in scenarios:
            label_f1 = macro_f1(y_label, oof_label[scenario])
            severity_accuracy = float(
                accuracy_score(
                    y_severity.to_numpy()[true_fault],
                    oof_severity[scenario][true_fault],
                )
            )
            score = raw_score(label_f1, severity_accuracy)
            per_class = f1_score(
                y_label,
                oof_label[scenario],
                labels=LABELS,
                average=None,
                zero_division=0,
            )
            row: dict[str, object] = {
                "cv_seed": cv_seed,
                "scenario": scenario,
                "macro_f1_oof": label_f1,
                "severity_accuracy_submission_oof": severity_accuracy,
                "raw_score_oof": score,
                "ml_points_oof": ml_points(score),
                "ood_overrides": ood_count[scenario],
            }
            row.update(
                {f"f1_{label}": float(value) for label, value in zip(LABELS, per_class, strict=True)}
            )
            run_rows.append(row)

    runs = pd.DataFrame(run_rows).sort_values(["scenario", "cv_seed"])
    summary = (
        runs.groupby("scenario", sort=False)
        .agg(
            seeds=("cv_seed", "nunique"),
            macro_f1_mean=("macro_f1_oof", "mean"),
            macro_f1_std=("macro_f1_oof", "std"),
            macro_f1_min=("macro_f1_oof", "min"),
            severity_mean=("severity_accuracy_submission_oof", "mean"),
            severity_std=("severity_accuracy_submission_oof", "std"),
            severity_min=("severity_accuracy_submission_oof", "min"),
            raw_score_mean=("raw_score_oof", "mean"),
            raw_score_std=("raw_score_oof", "std"),
            raw_score_min=("raw_score_oof", "min"),
            ml_points_mean=("ml_points_oof", "mean"),
            ml_points_min=("ml_points_oof", "min"),
            unknown_f1_mean=("f1_unknown", "mean"),
            ood_overrides_mean=("ood_overrides", "mean"),
        )
        .reset_index()
    )
    runs.to_csv(args.output_dir / "ml_polish_runs.csv", index=False)
    summary.to_csv(args.output_dir / "ml_polish_summary.csv", index=False)
    print("\nPolished architecture summary:")
    print(summary.to_string(index=False, float_format="{:.4f}".format))
    print(f"\nOutputs saved to: {args.output_dir.resolve()}")


if __name__ == "__main__":
    main()