Compressor/models/vibration/vibration_algorithm.py

#!/usr/bin/env python3
"""PKS vibration anomaly-segment algorithm (vibration half of the scanner).

Vibration-only reconstruction of ``backend/scan_pks_fault_evolution.py``.  The
ported functions are byte-for-byte equivalent to the original so the exported
``振动异常段.csv`` stays identical.  Shared query/IO/statistics helpers live in
:mod:`models.common`.

CLI
---
::

    python models/vibration/vibration_algorithm.py \
        --hourly-csv models/datas/hourly/pks_hourly_full.csv \
        --output-dir models/results/vibration
"""
from __future__ import annotations

import argparse
import sys
from collections import defaultdict
from datetime import timedelta
from pathlib import Path
from statistics import median

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import db, devices, io, paths, queries
from models.common.constants import (
    MIN_HOUR_COVERAGE,
    VIBRATION_BASELINE_DAYS,
    VIBRATION_MERGE_GAP_HOURS,
    VIBRATION_RESTART_GAP_HOURS,
)
from models.common.stats import median_value, robust_baseline


def vibration_hour_signals(rows: list[dict], configs: dict[int, dict]):
    signals = []
    grouped = defaultdict(list)
    for row in rows:
        if row["运行覆盖率"] >= MIN_HOUR_COVERAGE:
            grouped[row["批次"]].append(row)
    for batch, series in grouped.items():
        series.sort(key=lambda row: row["_dt"])
        config = configs[batch]
        valid = []
        last_valid_dt = None
        for row in series:
            if row["运行覆盖率"] < MIN_HOUR_COVERAGE:
                continue
            restart = last_valid_dt is None or row["_dt"] - last_valid_dt > timedelta(hours=VIBRATION_RESTART_GAP_HOURS)
            if restart:
                # The first complete running hour after a restart is a
                # transition hour. Keep older valid hours for the long-term
                # robust reference, but do not evaluate this transition hour.
                last_valid_dt = row["_dt"]
                continue
            baseline_start = row["_dt"] - timedelta(days=VIBRATION_BASELINE_DAYS)
            prior = [item for item in valid if item["_dt"] >= baseline_start]
            side_signals = []
            for side in ("coupling", "chain"):
                avg_center, avg_scale = robust_baseline([item[f"{side}_avg"] for item in prior])
                max_center, max_scale = robust_baseline([item[f"{side}_max"] for item in prior])
                current_avg = row[f"{side}_avg"]
                current_max = row[f"{side}_max"]
                if current_max is None:
                    continue
                avg_z = ((current_avg - avg_center) / avg_scale) if avg_center is not None and current_avg is not None else 0.0
                peak_z = (current_max - max_center) / max_scale if max_center is not None and max_scale else 0.0
                running = max(row["运行样本数"], 1)
                high_ratio = (row[f"{side}_high_count"] or 0) / running
                high_high_ratio = (row[f"{side}_high_high_count"] or 0) / running
                limits = config[side]
                local_prior = valid[:6] if len(valid) >= 6 else []
                local_avg = median_value([item[f"{side}_avg"] for item in local_prior])
                local_max = median_value([item[f"{side}_max"] for item in local_prior])
                local_rise = False
                severe = (
                    high_high_ratio >= 0.01
                    or current_max >= limits["high_high"]
                    or peak_z >= 8
                )
                abnormal = (
                    high_ratio >= 0.01
                    or current_max >= limits["high"]
                    or peak_z >= 5
                    or avg_z >= 4
                    or local_rise
                )
                mild = peak_z >= 4 or avg_z >= 3 or local_rise
                if severe or abnormal or mild:
                    side_signals.append({
                        "side": side,
                        "avg_z": avg_z,
                        "peak_z": peak_z,
                        "high_ratio": high_ratio,
                        "high_high_ratio": high_high_ratio,
                        "severe": severe,
                        "abnormal": abnormal,
                        "value": current_max,
                        "avg": current_avg,
                        "prior_max_center": max_center,
                        "prior_max_scale": max_scale,
                        "prior_avg_center": avg_center,
                    })
            if side_signals:
                strongest = max(side_signals, key=lambda item: (item["severe"], item["peak_z"], item["avg_z"]))
                signal_level = 3 if any(item["severe"] for item in side_signals) else 2 if any(item["abnormal"] for item in side_signals) else 1
                reasons = []
                for item in side_signals:
                    label = "联轴器端" if item["side"] == "coupling" else "链轮端"
                    reasons.append(
                        f"{label}均值偏离{item['avg_z']:.1f}倍、峰值偏离{item['peak_z']:.1f}倍、"
                        f"高报警比例{item['high_ratio']:.1%}"
                    )
                signals.append({
                    "批次": batch,
                    "机组": config["unit"],
                    "_dt": row["_dt"],
                    "_level": signal_level,
                    "_severe": signal_level == 3,
                    "_side_signals": side_signals,
                    "_strongest": strongest,
                    "_reason": "；".join(reasons),
                })
            valid.append(row)
            last_valid_dt = row["_dt"]
    return signals


def merge_vibration_signals(signals: list[dict], configs: dict[int, dict]):
    grouped = defaultdict(list)
    for signal in signals:
        grouped[signal["批次"]].append(signal)
    events = []
    for batch, series in grouped.items():
        series.sort(key=lambda item: item["_dt"])
        current = []
        for signal in series:
            if not current or signal["_dt"] - current[-1]["_dt"] <= timedelta(hours=VIBRATION_MERGE_GAP_HOURS):
                current.append(signal)
            else:
                events.append(build_vibration_event(current, configs[batch]))
                current = [signal]
        if current:
            events.append(build_vibration_event(current, configs[batch]))

    events.sort(key=lambda row: (queries.unit_sort_key(row["机组"]), row["开始小时"], row["批次"]))
    unit_indexes = defaultdict(int)
    for event in events:
        unit_indexes[event["机组"]] += 1
        event["周期编号"] = f'{event["机组"]}-V{unit_indexes[event["机组"]]:02d}'

    filtered = []
    for event in events:
        # A single post-start relative-rise signal is not sufficient evidence;
        # absolute severe signals remain eligible on their own.
        if event["_严重信号小时数"] > 0 or event["_触发小时数"] >= 2:
            filtered.append(event)
    return sorted(filtered, key=io.output_sort_key)


def vibration_score(duration_hours: int, valid_hours: int, avg_change: float, max_value: float, peak_z: float, high_hours: int, high_high_hours: int, config: dict, side: str):
    duration_score = min(40.0, max(0.0, (duration_hours - 1) / 12 * 40))
    avg_score = min(20.0, max(0.0, (avg_change - 0.10) / 0.40 * 20)) if avg_change >= 0.10 else 0.0
    high = config[side]["high"]
    high_high = config[side]["high_high"]
    if max_value < high:
        max_score = min(10.0, max(0.0, (max_value / max(high, 1e-6) - 0.8) / 0.2 * 10))
    elif max_value < high_high:
        max_score = 10.0 + (max_value - high) / max(high_high - high, 1e-6) * 10.0
    else:
        max_score = 20.0
    peak_score = min(20.0, max(0.0, (peak_z - 4.0) / 8.0 * 20.0)) if peak_z >= 4.0 else 0.0
    total = duration_score + avg_score + max_score + peak_score
    hard_severe = max_value >= high_high or high_high_hours > 0 or (peak_z >= 8.0 and valid_hours >= 2)
    if hard_severe or total >= 60.0:
        level = "严重异常"
    elif total >= 30.0:
        level = "异常"
    else:
        level = "轻微"
    return total, level, duration_score, avg_score, max_score, peak_score


def build_vibration_event(signals: list[dict], config: dict):
    start = signals[0]["_dt"]
    end = signals[-1]["_dt"]
    all_side = [item for signal in signals for item in signal["_side_signals"]]
    strongest = max(all_side, key=lambda item: (item["severe"], item["peak_z"], item["value"]))
    main_side = "联轴器端" if strongest["side"] == "coupling" else "链轮端"
    max_value = max(item["value"] for item in all_side if item["value"] is not None)
    max_peak_z = max(item["peak_z"] for item in all_side)
    high_hours = sum(
        1 for signal in signals
        if any(item["high_ratio"] > 0 for item in signal["_side_signals"])
    )
    high_high_hours = sum(
        1 for signal in signals
        if any(item["high_high_ratio"] > 0 for item in signal["_side_signals"])
    )
    severe_hours = sum(1 for signal in signals if signal["_severe"])
    duration_hours = int((end - start).total_seconds() / 3600) + 1
    running_days = len({signal["_dt"].date() for signal in signals})
    side_avg_changes = []
    for side in ("coupling", "chain"):
        values = [item["avg"] for item in all_side if item["side"] == side and item["avg"] is not None]
        prior_values = [item["prior_avg_center"] for item in all_side if item["side"] == side and item["prior_avg_center"] is not None]
        if values and prior_values and median(prior_values) != 0:
            side_avg_changes.append((median(values) / median(prior_values) - 1, side))
    avg_change, _ = max(side_avg_changes, key=lambda item: item[0], default=(0.0, strongest["side"]))
    peak_ratio = max_value / max(strongest["prior_max_center"], 1e-6)

    if max_value >= config[strongest["side"]]["high_high"] or severe_hours > 0 and (high_high_hours > 0 or max_peak_z >= 8):
        level = "严重异常"
    elif max_value >= config[strongest["side"]]["high"] and high_hours >= 2:
        level = "严重异常"
    elif max_value >= config[strongest["side"]]["high"] or max_peak_z >= 5 or avg_change >= 0.10:
        level = "异常"
    else:
        level = "轻微"

    if peak_ratio >= 1.5 and avg_change < 0.25:
        shape = "间歇冲击型"
    elif avg_change >= 0.10 and len(signals) >= 3:
        shape = "持续抬升型"
    else:
        shape = "混合型"

    peak_signal = max(signals, key=lambda signal: signal["_strongest"]["peak_z"])
    first_signal = signals[0]
    last_signal = signals[-1]
    first_value = first_signal["_strongest"]["value"]
    last_value = last_signal["_strongest"]["value"]
    baseline = strongest["prior_max_center"]
    scale = max(strongest.get("prior_max_scale", 0.0), 0.0)
    normal_band = f"{baseline - 2 * scale:.3f}～{baseline + 2 * scale:.3f} mm/s" if baseline is not None else ""
    consecutive = 0
    longest_consecutive = 0
    for signal in signals:
        if signal["_level"] >= 2:
            consecutive += 1
            longest_consecutive = max(longest_consecutive, consecutive)
        else:
            consecutive = 0
    score, level, duration_score, avg_score, max_score, peak_score = vibration_score(
        duration_hours, len(signals), avg_change, max_value, max_peak_z,
        high_hours, high_high_hours, config, strongest["side"]
    )
    confirmed = "是" if level != "轻微" else "否"
    if shape == "持续抬升型":
        basis = f"{main_side}振动均值持续抬升，阶段内{longest_consecutive or len(signals)}个有效运行小时触发"
    elif shape == "间歇冲击型":
        basis = f"{main_side}重复出现振动冲击，阶段内{len(signals)}个有效运行小时触发"
    else:
        basis = f"{main_side}振动峰值或均值偏离近期基线，阶段内{len(signals)}个有效运行小时触发"
    if high_high_hours:
        basis += f"；高高报警小时{high_high_hours}个"
    elif high_hours:
        basis += f"；高报警小时{high_hours}个"
    basis += f"；最大振动值{max_value:.3f}mm/s，峰值高于近期基线{max_peak_z:.2f}个稳健尺度（峰值/基线{peak_ratio:.2f}倍）；综合评分{score:.1f}分"
    return {
        "机组": config["unit"],
        "批次": config["batch"],
        "周期编号": "",
        "振动阶段": level,
        "开始小时": start.strftime("%Y-%m-%d %H:%M:%S"),
        "结束小时": end.strftime("%Y-%m-%d %H:%M:%S"),
        "持续自然小时数": duration_hours,
        "有效运行小时数": len(signals),
        "近期基线振动": baseline,
        "正常波动带": normal_band,
        "阶段开始小时振动": first_value,
        "阶段结束小时振动": last_value,
        "阶段最大振动值": max_value,
        "阶段开始相对近期基线变化": (first_value / baseline - 1) if baseline else None,
        "阶段结束相对近期基线变化": (last_value / baseline - 1) if baseline else None,
        "阶段最大连续异常有效小时数": longest_consecutive,
        "阶段综合评分": score,
        "持续时间得分": duration_score,
        "平均振动得分": avg_score,
        "最大振动值得分": max_score,
        "峰值偏离得分": peak_score,
        "最强异常侧": main_side,
        "最强侧近期峰值基线": baseline,
        "最强侧近期稳健尺度": scale,
        "阶段最大峰值/基线比例": peak_ratio,
        "阶段高报警小时数": high_hours,
        "阶段高高报警小时数": high_high_hours,
        "_严重信号小时数": severe_hours,
        "_触发小时数": len(signals),
        "是否形成确认等级": confirmed,
        "阶段触发依据": basis,
    }


VIBRATION_FIELDS = [
    "机组", "批次", "周期编号", "振动阶段", "开始小时", "结束小时", "持续自然小时数", "有效运行小时数",
    "近期基线振动", "正常波动带", "阶段开始小时振动", "阶段结束小时振动", "阶段最大振动值",
    "阶段开始相对近期基线变化", "阶段结束相对近期基线变化", "阶段最大连续异常有效小时数",
    "阶段高报警小时数", "阶段高高报警小时数", "是否形成确认等级", "阶段触发依据",
    "阶段综合评分", "持续时间得分", "平均振动得分", "最大振动值得分", "峰值偏离得分",
    "最强异常侧", "最强侧近期峰值基线", "最强侧近期稳健尺度", "阶段最大峰值/基线比例",
]


def run(rows: list[dict], configs: dict[int, dict]):
    """Return vibration anomaly events."""
    signals = vibration_hour_signals(rows, configs)
    return merge_vibration_signals(signals, configs)


def write_results(output_dir: Path, events, extra_fields=()):
    output_dir = Path(output_dir)
    fields = [field for field in VIBRATION_FIELDS if field not in extra_fields] + list(extra_fields)
    io.write_csv(output_dir / "振动异常段.csv", io.clean_rows(events), fields, {field: field for field in fields})


def load_hourly(args, selected_batches, unit_by_batch, configs):
    if args.hourly_csv:
        selected_set = set(selected_batches)
        return [row for row in queries.load_hourly_csv(Path(args.hourly_csv)) if row["批次"] in selected_set]
    start, end = queries.parse_dt(args.start), queries.parse_dt(args.end)
    return queries.fetch_all_hourly(start, end, args.chunk_days, unit_by_batch, configs, selected_batches)


def main() -> None:
    parser = argparse.ArgumentParser(description="PKS 振动异常段算法（独立于压力）")
    parser.add_argument("--device", nargs="+", default=list(devices.DEVICE_CODES), help="压缩机编号，如 7# 8# 9#")
    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("--hourly-csv", default=None)
    parser.add_argument("--output-dir", default=str(paths.VIBRATION_RESULTS))
    args = parser.parse_args()
    devices_selected = devices.normalize_devices(args.device)
    device_batches = devices.load_device_batches(devices_selected)
    unit_by_batch = devices.batch_to_unit(device_batches)
    configs = db.load_alarm_config(unit_by_batch)
    selected_batches = sorted(unit_by_batch)
    rows = load_hourly(args, selected_batches, unit_by_batch, configs)
    events = run(rows, configs)
    write_results(Path(args.output_dir), events)
    print(f"设备：{devices_selected}")
    print(f"小时特征：{len(rows):,} 行")
    print(f"振动异常段：{len(events):,} 段")
    print(f"输出目录：{args.output_dir}")


if __name__ == "__main__":
    main()