Compressor/backend/model_assistant_generalized/train_model.py

#!/usr/bin/env python3
"""Train one XGBoost PKS assistant model without reading the test set."""
from __future__ import annotations

import argparse
import importlib
import json
from collections import defaultdict
from pathlib import Path

import joblib
import numpy as np
from sklearn.impute import SimpleImputer
from sklearn.metrics import average_precision_score, precision_recall_curve
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from xgboost import XGBClassifier
from xgboost.callback import TrainingCallback


def feature_module(model_type: str):
    """Pressure and vibration use independent feature modules."""
    return importlib.import_module(
        "generalized_features" if model_type == "pressure" else "generalized_features_vibration"
    )


class ProgressCallback(TrainingCallback):
    """Print round, learning rate, and train/validation loss for each boosting round."""

    def __init__(self, rounds: int, learning_rate: float) -> None:
        self.rounds = rounds
        self.learning_rate = learning_rate

    def after_iteration(self, model, epoch: int, evals_log):
        losses = []
        for dataset, metrics in evals_log.items():
            for name, values in metrics.items():
                if values:
                    losses.append(f"{dataset}-{name}={values[-1]:.5f}")
        print(
            f"[round {epoch + 1:>4}/{self.rounds}] lr={self.learning_rate:.4f} " + " ".join(losses),
            flush=True,
        )
        return False


def load_samples(directory: Path) -> list[dict]:
    return [json.loads(path.read_text(encoding="utf-8")) for path in sorted(directory.glob("*.json"))]


def matrix(feature_dicts: list[dict], names: list[str]) -> np.ndarray:
    return np.asarray([[item.get(name) for name in names] for item in feature_dicts], dtype=float)


def stratified_normal_sample(samples: list[dict], max_ratio: int, seed: int) -> list[dict]:
    positives = [sample for sample in samples if sample["label"] == 1]
    normals = [sample for sample in samples if sample["label"] == 0]
    limit = len(positives) * max_ratio
    if len(normals) <= limit:
        return samples
    buckets = defaultdict(list)
    for sample in normals:
        month = sample["start_hour"][:7]
        buckets[(sample["unit"], sample["batch"], month)].append(sample)
    rng = np.random.default_rng(seed)
    selected = []
    # Round-robin quotas preserve unit/month coverage rather than favoring the
    # longest uninterrupted running cycle.
    keys = sorted(buckets)
    quota, remainder = divmod(limit, len(keys))
    for index, key in enumerate(keys):
        bucket = buckets[key]
        take = min(len(bucket), quota + (1 if index < remainder else 0))
        selected.extend(bucket[index] for index in rng.choice(len(bucket), size=take, replace=False))
    if len(selected) < limit:
        leftovers = [sample for sample in normals if sample not in selected]
        need = min(limit - len(selected), len(leftovers))
        selected.extend(leftovers[index] for index in rng.choice(len(leftovers), size=need, replace=False))
    return positives + selected


def choose_threshold(y: np.ndarray, probabilities: np.ndarray) -> tuple[float, dict]:
    precision, recall, thresholds = precision_recall_curve(y, probabilities)
    f1 = 2 * precision * recall / np.maximum(precision + recall, 1e-12)
    index = int(np.nanargmax(f1[:-1])) if len(thresholds) else 0
    threshold = float(thresholds[index]) if len(thresholds) else 0.5
    return threshold, {
        "threshold_selection": "max_validation_f1",
        "threshold": threshold,
        "validation_precision": float(precision[index]),
        "validation_recall": float(recall[index]),
        "validation_f1": float(f1[index]),
        "validation_pr_auc": float(average_precision_score(y, probabilities)),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="训练 PKS 压力或振动模型，不读取测试集")
    parser.add_argument("model_type", choices=("pressure", "vibration"))
    parser.add_argument("--datasets-dir", default=str(Path(__file__).resolve().parent / "datasets"))
    parser.add_argument("--models-dir", default=str(Path(__file__).resolve().parent / "models"))
    sampling = parser.add_mutually_exclusive_group()
    sampling.add_argument("--normal-ratio", type=int, default=None)
    sampling.add_argument("--all-normal", action="store_true", help="使用全部正常训练样本，不下采样、不复制")
    parser.add_argument("--seed", type=int, default=20260921)
    parser.add_argument("--rounds", type=int, default=250)
    parser.add_argument("--early-stopping-rounds", type=int, default=30)
    args = parser.parse_args()
    if args.normal_ratio is not None and args.normal_ratio < 1:
        parser.error("--normal-ratio 必须大于等于 1")
    if args.rounds < 1 or args.early_stopping_rounds < 1:
        parser.error("--rounds 和 --early-stopping-rounds 必须大于等于 1")
    ratio = args.normal_ratio if args.normal_ratio is not None else (4 if args.model_type == "pressure" else 5)
    gf = feature_module(args.model_type)
    dataset = Path(args.datasets_dir) / args.model_type
    train_all = load_samples(dataset / "train")
    validation = load_samples(dataset / "validation")
    if not train_all or not validation:
        parser.error("训练或验证样本为空")
    if not any(sample["label"] == 1 for sample in train_all) or not any(sample["label"] == 1 for sample in validation):
        parser.error("训练和验证集都必须包含异常样本")
    train = train_all if args.all_normal else stratified_normal_sample(train_all, ratio, args.seed)
    baselines = gf.compute_unit_baselines(train_all)
    names = gf.feature_names()
    x_train, y_train = matrix(gf.feature_matrix(train, baselines), names), np.asarray([sample["label"] for sample in train])
    x_validation, y_validation = matrix(gf.feature_matrix(validation, baselines), names), np.asarray([sample["label"] for sample in validation])
    imputer = SimpleImputer(strategy="median")
    x_train = imputer.fit_transform(x_train)
    x_validation = imputer.transform(x_validation)
    scale_pos_weight = float(np.sum(y_train == 0) / max(np.sum(y_train == 1), 1))
    learning_rate = 0.05
    print(
        f"training: rounds={args.rounds}, learning_rate={learning_rate:.4f}, "
        f"early_stopping_rounds={args.early_stopping_rounds}, samples={len(train):,}",
        flush=True,
    )
    model = XGBClassifier(
        n_estimators=args.rounds, max_depth=3, learning_rate=learning_rate,
        subsample=0.8, colsample_bytree=0.8, min_child_weight=3,
        objective="binary:logistic", eval_metric="logloss", scale_pos_weight=scale_pos_weight,
        random_state=args.seed, n_jobs=4, early_stopping_rounds=args.early_stopping_rounds,
        callbacks=[ProgressCallback(args.rounds, learning_rate)],
    )
    eval_set = [(x_train, y_train), (x_validation, y_validation)]
    model.fit(x_train, y_train, eval_set=eval_set, verbose=False)
    rounds_run = len(model.evals_result()["validation_0"]["logloss"])
    stopped_early = rounds_run < args.rounds
    print(
        f"training complete: rounds_run={rounds_run}/{args.rounds}, "
        f"best_iteration={model.best_iteration + 1}, best_validation_logloss={float(model.best_score):.5f}, "
        f"early_stopped={'yes' if stopped_early else 'no'}",
        flush=True,
    )
    validation_probabilities = model.predict_proba(x_validation)[:, 1]
    threshold, validation_metrics = choose_threshold(y_validation, validation_probabilities)
    similarity = Pipeline([("scale", StandardScaler()), ("knn", KNeighborsClassifier(n_neighbors=min(5, len(train)), weights="distance"))])
    similarity.fit(x_train, y_train)
    models = Path(args.models_dir)
    models.mkdir(parents=True, exist_ok=True)
    # The progress callback is defined in this script (module __main__); drop it
    # before serialization so the artifact can be loaded by other entry points.
    model.callbacks = []
    joblib.dump(
        {"model": model, "imputer": imputer, "knn": similarity, "feature_names": names, "unit_baselines": baselines},
        models / f"{args.model_type}_model.joblib",
    )
    (models / f"{args.model_type}_unit_baselines.json").write_text(
        json.dumps(baselines, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    metadata = {
        "model_type": args.model_type, "algorithm": "XGBoost binary classifier + KNN similarity",
        "feature_version": "generalized_cycle_relative_v2",
        "feature_names": names,
        "dropped_absolute_features": list(gf.ABSOLUTE_FEATURES),
        "unit_baselines": baselines,
        "normal_sampling_strategy": "all_normal" if args.all_normal else "stratified_max_ratio",
        "normal_to_anomaly_training_ratio_limit": None if args.all_normal else ratio,
        "actual_normal_to_anomaly_training_ratio": float(np.sum(y_train == 0) / max(np.sum(y_train == 1), 1)),
        "training_samples_before_sampling": len(train_all), "training_samples_after_sampling": len(train),
        "training_normal": int(np.sum(y_train == 0)), "training_anomaly": int(np.sum(y_train == 1)),
        "scale_pos_weight": scale_pos_weight, "validation_samples": len(validation),
        "test_set_read": False, **validation_metrics,
        "rounds_requested": args.rounds, "early_stopping_rounds_requested": args.early_stopping_rounds,
        "rounds_run": rounds_run, "early_stopped": stopped_early,
        "best_iteration": int(model.best_iteration), "best_score": float(model.best_score),
    }
    (models / f"{args.model_type}_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()