Compressor/models/vibration/build_datasets.py

#!/usr/bin/env python3
"""Build vibration model samples from the reconstructed vibration algorithm output."""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import paths, queries, windowing
from models.common.constants import VIBRATION_BASELINE_DAYS
from models.vibration import features_vibration


def vibration_features(series, end_index, window):
    features = {
        "running_coverage_mean": windowing.mean([row["运行覆盖率"] for row in window]),
        "running_coverage_min": min(row["运行覆盖率"] for row in window),
    }
    for side in ("coupling", "chain"):
        averages = [row[f"{side}_avg"] for row in window]
        peaks = [row[f"{side}_max"] for row in window]
        avg_center, avg_scale = windowing.previous_baseline(series, end_index, f"{side}_avg", VIBRATION_BASELINE_DAYS * 24)
        peak_center, peak_scale = windowing.previous_baseline(series, end_index, f"{side}_max", VIBRATION_BASELINE_DAYS * 24)
        current_avg, current_peak = averages[-1], peaks[-1]
        features.update({
            f"{side}_avg_mean": windowing.mean(averages), f"{side}_avg_max": max(averages), f"{side}_avg_range": windowing.value_range(averages),
            f"{side}_avg_slope": windowing.slope(averages), f"{side}_peak_mean": windowing.mean(peaks), f"{side}_peak_max": max(peaks),
            f"{side}_peak_range": windowing.value_range(peaks), f"{side}_peak_slope": windowing.slope(peaks),
            f"{side}_avg_baseline_deviation": ((current_avg - avg_center) / avg_scale) if avg_center is not None and avg_scale else None,
            f"{side}_peak_baseline_deviation": ((current_peak - peak_center) / peak_scale) if peak_center is not None and peak_scale else None,
            f"{side}_avg_rel_last": (current_avg / avg_center - 1) if avg_center else None,
            f"{side}_peak_rel_last": (current_peak / peak_center - 1) if peak_center else None,
            f"{side}_avg_rel_mean": (windowing.mean(averages) / avg_center - 1) if avg_center else None,
            f"{side}_peak_rel_max": (max(peaks) / peak_center - 1) if peak_center else None,
            f"{side}_high_ratio_mean": windowing.mean([windowing.ratio(row[f"{side}_high_count"], row["运行样本数"]) for row in window]),
            f"{side}_high_high_ratio_mean": windowing.mean([windowing.ratio(row[f"{side}_high_high_count"], row["运行样本数"]) for row in window]),
        })
    features["both_sides_rising"] = int(features["coupling_avg_slope"] > 0 and features["chain_avg_slope"] > 0)
    return features


def main() -> None:
    parser = argparse.ArgumentParser(description="构建振动模型样本（窗口可配置）")
    parser.add_argument("--start", default="2025-01-01 00:00:00")
    parser.add_argument("--end", default="2026-09-01 00:00:00")
    parser.add_argument("--chunk-days", type=int, default=14)
    parser.add_argument("--window-hours", type=int, default=3)
    parser.add_argument("--output-dir", default=str(paths.VIBRATION_DATAS))
    parser.add_argument("--hourly-csv", default=str(paths.HOURLY_FULL))
    parser.add_argument("--units", nargs="+", default=None)
    args = parser.parse_args()

    output = Path(args.output_dir)
    if output.exists():
        shutil.rmtree(output)
    output.mkdir(parents=True)

    hourly_path = Path(args.hourly_csv)
    if hourly_path.is_file():
        print(f"复用小时中间表：{hourly_path}")
        rows = queries.load_hourly_csv(hourly_path)
    else:
        from models.common import db

        batch_to_unit = db.load_pks_batches()
        selected = sorted(batch_to_unit)
        configs = db.load_alarm_config({batch: batch_to_unit[batch] for batch in selected})
        start, end = queries.parse_dt(args.start), queries.parse_dt(args.end)
        rows = queries.fetch_all_hourly(start, end, args.chunk_days, batch_to_unit, configs, selected)

    scoped = [row for row in rows if not args.units or row["机组"] in set(args.units)]
    windowing.write_hourly(output / "hourly.csv", scoped)
    stages = windowing.stage_ranges(paths.VIBRATION_RESULTS / "振动异常段.csv", "振动阶段")
    if args.units:
        stages = [stage for stage in stages if stage["unit"] in set(args.units)]
    samples = windowing.build_samples(scoped, "vibration", stages, args.window_hours, vibration_features)
    counts = windowing.write_samples(output, samples)
    print(f"vibration samples: {dict(counts)}")
    print(f"小时中间数据：{len(scoped):,} 行")
    print(f"样本输出目录：{output}")


if __name__ == "__main__":
    main()