"""Generic XGBoost training / evaluation / inference for both model types. Ports the original generalized model pipeline's training, evaluation and inference logic. The only model-specific input is the feature module (``features_pressure`` or ``features_vibration``). """ from __future__ import annotations import csv 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, confusion_matrix, precision_recall_curve, precision_recall_fscore_support from sklearn.neighbors import KNeighborsClassifier, NearestNeighbors from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler from xgboost import XGBClassifier from xgboost.callback import TrainingCallback def load_samples(directory: Path) -> list[dict]: return [json.loads(path.read_text(encoding="utf-8")) for path in sorted(Path(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) class ProgressCallback(TrainingCallback): 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 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 = [] 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 train(model_type: str, feature_module, dataset_dir: Path, models_dir: Path, seed: int, rounds: int, early_stopping_rounds: int, all_normal: bool, normal_ratio: int | None) -> dict: ratio = normal_ratio if normal_ratio is not None else (4 if model_type == "pressure" else 5) dataset = Path(dataset_dir) train_all = load_samples(dataset / "train") validation = load_samples(dataset / "validation") if not train_all or not validation: raise RuntimeError("训练或验证样本为空") if not any(sample["label"] == 1 for sample in train_all) or not any(sample["label"] == 1 for sample in validation): raise RuntimeError("训练和验证集都必须包含异常样本") train = train_all if all_normal else stratified_normal_sample(train_all, ratio, seed) baselines = feature_module.compute_unit_baselines(train_all) names = feature_module.feature_names() x_train, y_train = matrix(feature_module.feature_matrix(train, baselines), names), np.asarray([sample["label"] for sample in train]) x_validation, y_validation = matrix(feature_module.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={rounds}, learning_rate={learning_rate:.4f}, early_stopping_rounds={early_stopping_rounds}, samples={len(train):,}", flush=True) model = XGBClassifier( n_estimators=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=seed, n_jobs=4, early_stopping_rounds=early_stopping_rounds, callbacks=[ProgressCallback(rounds, learning_rate)], ) model.fit(x_train, y_train, eval_set=[(x_train, y_train), (x_validation, y_validation)], verbose=False) rounds_run = len(model.evals_result()["validation_0"]["logloss"]) stopped_early = rounds_run < rounds print(f"training complete: rounds_run={rounds_run}/{rounds}, best_iteration={model.best_iteration + 1}, best_validation_logloss={float(model.best_score):.5f}, 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(models_dir) models.mkdir(parents=True, exist_ok=True) model.callbacks = [] joblib.dump( {"model": model, "imputer": imputer, "knn": similarity, "feature_names": names, "unit_baselines": baselines}, models / f"{model_type}_model.joblib", ) (models / f"{model_type}_unit_baselines.json").write_text(json.dumps(baselines, ensure_ascii=False, indent=2), encoding="utf-8") metadata = { "model_type": model_type, "algorithm": "XGBoost binary classifier + KNN similarity", "feature_version": "generalized_cycle_relative_v2", "feature_names": names, "dropped_absolute_features": list(feature_module.ABSOLUTE_FEATURES), "unit_baselines": baselines, "normal_sampling_strategy": "all_normal" if all_normal else "stratified_max_ratio", "normal_to_anomaly_training_ratio_limit": None if 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": rounds, "early_stopping_rounds_requested": 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"{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)) return metadata def metrics(y: np.ndarray, predicted: np.ndarray, probabilities: np.ndarray) -> dict: labels = [0, 1] per_precision, per_recall, per_f1, per_support = precision_recall_fscore_support(y, predicted, labels=labels, zero_division=0) tn, fp, fn, tp = confusion_matrix(y, predicted, labels=labels).ravel() samples = int(len(y)) normal_support, anomaly_support = int(per_support[0]), int(per_support[1]) specificity = float(tn / (tn + fp)) if (tn + fp) else 0.0 accuracy = float((tp + tn) / samples) if samples else 0.0 return { "samples": samples, "normal_samples": normal_support, "anomaly_samples": anomaly_support, "tp": int(tp), "fp": int(fp), "fn": int(fn), "tn": int(tn), "precision": float(per_precision[1]), "recall": float(per_recall[1]), "f1": float(per_f1[1]), "pr_auc": float(average_precision_score(y, probabilities)) if len(set(y)) == 2 else None, "specificity": specificity, "accuracy": accuracy, "balanced_accuracy": float((per_recall[0] + per_recall[1]) / 2), "macro_f1": float((per_f1[0] + per_f1[1]) / 2), "weighted_f1": float((per_f1[0] * normal_support + per_f1[1] * anomaly_support) / samples) if samples else 0.0, "class_metrics": { "normal": {"precision": float(per_precision[0]), "recall": float(per_recall[0]), "f1": float(per_f1[0]), "support": normal_support}, "anomaly": {"precision": float(per_precision[1]), "recall": float(per_recall[1]), "f1": float(per_f1[1]), "support": anomaly_support}, }, "confusion_matrix_counts": {"tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp)}, "confusion_matrix_normalized": { "normal": {"predicted_normal": float(tn / normal_support) if normal_support else 0.0, "predicted_anomaly": float(fp / normal_support) if normal_support else 0.0}, "anomaly": {"predicted_normal": float(fn / anomaly_support) if anomaly_support else 0.0, "predicted_anomaly": float(tp / anomaly_support) if anomaly_support else 0.0}, }, } def event_metrics(samples: list[dict], predicted: np.ndarray) -> dict: groups = defaultdict(list) for sample, value in zip(samples, predicted): groups[sample["event_group_id"]].append((sample, int(value))) positives = [values for values in groups.values() if values[0][0]["label"] == 1] negatives = [values for values in groups.values() if values[0][0]["label"] == 0] tp = sum(any(value for _, value in values) for values in positives) fn = len(positives) - tp fp = sum(any(value for _, value in values) for values in negatives) return { "anomaly_events": len(positives), "normal_events": len(negatives), "detected_anomaly_events": tp, "missed_anomaly_events": fn, "false_positive_normal_events": fp, "event_recall": tp / len(positives) if positives else None, } def evaluate(model_type: str, feature_module, dataset_dir: Path, models_dir: Path, output_dir: Path) -> dict: bundle = joblib.load(Path(models_dir) / f"{model_type}_model.joblib") metadata = json.loads((Path(models_dir) / f"{model_type}_metadata.json").read_text(encoding="utf-8")) samples = load_samples(Path(dataset_dir) / "test") if not samples or not any(sample["label"] == 1 for sample in samples): raise RuntimeError("封存测试集必须包含至少一个异常样本") x = bundle["imputer"].transform(matrix(feature_module.feature_matrix(samples, bundle["unit_baselines"]), bundle["feature_names"])) y = np.asarray([sample["label"] for sample in samples]) probabilities = bundle["model"].predict_proba(x)[:, 1] predicted = (probabilities >= metadata["threshold"]).astype(int) distances, neighbors = bundle["knn"].named_steps["knn"].kneighbors(bundle["knn"].named_steps["scale"].transform(x), n_neighbors=1) rows = [] for index, sample in enumerate(samples): rows.append({ "sample_id": sample["sample_id"], "unit": sample["unit"], "batch": sample["batch"], "event_group_id": sample["event_group_id"], "start_hour": sample["start_hour"], "end_hour": sample["end_hour"], "true_label": sample["label"], "anomaly_probability": probabilities[index], "predicted_label": int(predicted[index]), "nearest_training_sample_index": int(neighbors[index][0]), "nearest_training_distance": float(distances[index][0]), "nearest_training_similarity": float(1 / (1 + distances[index][0])), }) by_unit = {} for unit in sorted({sample["unit"] for sample in samples}): indices = [index for index, sample in enumerate(samples) if sample["unit"] == unit] by_unit[unit] = metrics(y[indices], predicted[indices], probabilities[indices]) report = { "model_type": model_type, "evaluation_set": "sealed_test", "threshold": metadata["threshold"], "window_metrics": metrics(y, predicted, probabilities), "event_metrics": event_metrics(samples, predicted), "by_unit": by_unit, "warning": "Labels originate from current scanner rules unless manually updated; metrics measure held-out rule-label agreement, not confirmed physical-fault diagnosis.", } output = Path(output_dir) output.mkdir(parents=True, exist_ok=True) with (output / f"{model_type}_test_predictions.csv").open("w", encoding="utf-8-sig", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=list(rows[0])) writer.writeheader() writer.writerows(rows) (output / f"{model_type}_test_evaluation.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") return report def combined_result(rule_label: int, probability: float, anomaly_similarity: float, threshold: float) -> str: model_support = probability >= threshold or anomaly_similarity >= 0.70 if rule_label == 1 and model_support: return "模型辅助确认异常" if rule_label == 1: return "保留规则异常，模型意见不一致，人工复核" if model_support: return "规则正常但模型疑似异常，继续观察" return "正常" def infer(model_type: str, feature_module, dataset_dir: Path, models_dir: Path, output_dir: Path) -> dict: dataset = Path(dataset_dir) bundle = joblib.load(Path(models_dir) / f"{model_type}_model.joblib") metadata = json.loads((Path(models_dir) / f"{model_type}_metadata.json").read_text(encoding="utf-8")) samples = load_samples(dataset / "candidates") training = load_samples(dataset / "train") baselines = bundle["unit_baselines"] x = bundle["imputer"].transform(matrix(feature_module.feature_matrix(samples, baselines), bundle["feature_names"])) x_train = bundle["imputer"].transform(matrix(feature_module.feature_matrix(training, baselines), bundle["feature_names"])) probabilities = bundle["model"].predict_proba(x)[:, 1] scaled = bundle["knn"].named_steps["scale"].transform(x) train_scaled = bundle["knn"].named_steps["scale"].transform(x_train) anomaly_indexes = [index for index, sample in enumerate(training) if sample["label"] == 1] if not anomaly_indexes: raise RuntimeError("训练集没有异常样本，无法建立异常相似性索引") anomaly_neighbors = NearestNeighbors(n_neighbors=1).fit(train_scaled[anomaly_indexes]) distances, nearest_indexes = anomaly_neighbors.kneighbors(scaled) rows = [] for index, sample in enumerate(samples): nearest = training[anomaly_indexes[int(nearest_indexes[index][0])]] anomaly_similarity = float(1 / (1 + distances[index][0])) rule_label = int(sample["label"]) rows.append({ "机组": sample["unit"], "批次": sample["batch"], "窗口开始小时": sample["start_hour"], "窗口结束小时": sample["end_hour"], "样本编号": sample["sample_id"], "现有规则结果": "异常" if rule_label else "正常", "规则阶段编号": sample.get("rule_stage_id") or "", "模型异常概率": round(float(probabilities[index]), 6), "最相似历史异常样本": nearest["sample_id"], "历史异常相似度": round(anomaly_similarity, 6), "联合判断": combined_result(rule_label, probabilities[index], anomaly_similarity, metadata["threshold"]), }) rows.sort(key=lambda row: (int(row["机组"][0]), row["窗口结束小时"], row["批次"])) output = Path(output_dir) output.mkdir(parents=True, exist_ok=True) name = "压力模型辅助确认.csv" if model_type == "pressure" else "振动模型辅助确认.csv" with (output / name).open("w", encoding="utf-8-sig", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=list(rows[0])) writer.writeheader() writer.writerows(rows) return {"rows": len(rows), "output": str(output / name)} __all__ = [ "load_samples", "matrix", "train", "evaluate", "infer", "metrics", "event_metrics", "choose_threshold", ]