Compressor/_tmp_cylinder_vibration/plots_actual.py

"""第二轮图：多缸并排的实际振动特征 + 首周期波形。"""
from __future__ import annotations

from datetime import timedelta

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from . import config, data

LEVEL_COLORS = {"严重": "#e74c3c", "异常": "#e67e22", "轻微": "#f1c40f"}


def _gap_aware(ax, x, y, gap_minutes: int = 60, **kwargs) -> None:
    x = pd.Series(pd.to_datetime(x)).reset_index(drop=True)
    y = pd.Series(np.asarray(y, dtype=float)).reset_index(drop=True)
    y = y.mask(x.diff() > pd.Timedelta(minutes=gap_minutes))
    ax.plot(x, y, **kwargs)


def _ratio_features(frame: pd.DataFrame, baseline: dict) -> pd.DataFrame:
    frame = frame.copy()
    frame["rms_ratio"] = np.nan
    frame["kurt_ratio"] = np.nan
    for cylinder, group in frame.groupby("cylinder", sort=False):
        meta = baseline.get(cylinder, {})
        rms_median = float(meta.get("rms_log", [0.0, 1.0])[0])
        kurt_median = float(meta.get("kurt", [1.0, 1.0])[0])
        frame.loc[group.index, "rms_ratio"] = np.power(10.0, group["rms_log"] - rms_median)
        if kurt_median > 0:
            frame.loc[group.index, "kurt_ratio"] = group["kurt"] / kurt_median
    return frame


def _draw_intervals(ax, intervals: pd.DataFrame, cylinder: str, start: pd.Timestamp, end: pd.Timestamp) -> None:
    if intervals.empty:
        return
    part = intervals[
        (intervals["cylinder"] == cylinder)
        & (intervals["end_time"] >= start)
        & (intervals["start_time"] <= end)
    ]
    for row in part.itertuples(index=False):
        color = LEVEL_COLORS.get(row.level, "#95a5a6")
        ax.axvspan(row.start_time, row.end_time, color=color, alpha=0.13)


def plot_feature_grid(frame: pd.DataFrame, intervals: pd.DataFrame, baseline: dict, event: dict) -> None:
    config.PLOT_DIR.mkdir(parents=True, exist_ok=True)
    cylinder = event["cylinder"]
    center = pd.Timestamp(event["date"])
    start = center - pd.Timedelta(days=10)
    end = center + pd.Timedelta(days=3)
    part = frame[(frame["sample_time"] >= start) & (frame["sample_time"] <= end)]
    if part.empty:
        return
    part = _ratio_features(part, baseline)

    for metric, threshold, ylabel in (
        ("rms_ratio", 1.5, "RMS / baseline"),
        ("kurt_ratio", 2.0, "kurtosis / baseline"),
    ):
        fig, axes = plt.subplots(2, 3, figsize=(15, 7), sharex=True)
        for ax, name in zip(axes.ravel(), config.CYLINDERS):
            group = part[part["cylinder"] == name].sort_values("sample_time")
            ax.set_title(f"cyl{config.CYLINDER_TO_INDEX[name]}", fontsize=9)
            ax.axvline(center, color="#2c3e50", linewidth=1.0)
            if group.empty:
                ax.text(0.5, 0.5, "no data in event window", ha="center", va="center", fontsize=8, color="#7f8c8d")
                ax.set_xticks([])
                ax.set_yticks([])
                continue
            _gap_aware(ax, group["sample_time"], group[metric], color="#c0392b" if name == cylinder else "#2980b9", linewidth=1.0)
            ax.axhline(threshold, color="#7f8c8d", linestyle="--", linewidth=0.9)
            _draw_intervals(ax, intervals, name, start, end)
            ax.set_ylabel(ylabel if name == cylinder else "", fontsize=8)
            ax.grid(alpha=0.2)
            ax.tick_params(labelsize=7)
        fig.suptitle(f"{event['event_id']} cylinder#{config.CYLINDER_TO_INDEX[cylinder]} - {metric} (gaps = no running data)")
        fig.tight_layout()
        path = config.PLOT_DIR / f"features_{metric}_{event['event_id']}.png"
        fig.savefig(path, dpi=130)
        plt.close(fig)
        print(f"[plot] {path}", flush=True)


def _nearest_row(frame: pd.DataFrame, cylinder: str, target: pd.Timestamp, tolerance: timedelta) -> pd.Series | None:
    group = frame[frame["cylinder"] == cylinder]
    if group.empty:
        return None
    distance = (group["sample_time"] - target).abs()
    index = distance.idxmin()
    if distance.loc[index] > tolerance:
        return None
    return group.loc[index]


def plot_waveform_grid(frame: pd.DataFrame, intervals: pd.DataFrame, event: dict) -> None:
    cylinder = event["cylinder"]
    if intervals.empty:
        return
    candidates = intervals[intervals["cylinder"] == cylinder]
    if candidates.empty:
        return
    seed = candidates.loc[candidates["max_score"].idxmax()]
    seed_time = pd.Timestamp(seed["max_time"])
    baseline_time = seed_time - pd.Timedelta(hours=24)

    seed_rows: dict[str, pd.Series] = {}
    baseline_rows: dict[str, pd.Series] = {}
    for name in config.CYLINDERS:
        row = _nearest_row(frame, name, seed_time, timedelta(hours=2))
        if row is not None:
            seed_rows[name] = row
        base = _nearest_row(frame, name, baseline_time, timedelta(hours=6))
        if base is None:
            group = frame[(frame["cylinder"] == name) & (frame["sample_time"] < seed_time)]
            base = group.sort_values("sample_time").tail(1).squeeze() if not group.empty else None
        if base is not None:
            baseline_rows[name] = base

    seed_ids = [int(row["id"]) for row in seed_rows.values()]
    baseline_ids = [int(row["id"]) for row in baseline_rows.values()]
    seed_samples = data.fetch_samples(seed_ids)
    baseline_samples = data.fetch_samples(baseline_ids)

    fig, axes = plt.subplots(2, 3, figsize=(15, 7), sharex=True)
    for ax, name in zip(axes.ravel(), config.CYLINDERS):
        row = seed_rows.get(name)
        if row is None:
            ax.text(0.5, 0.5, "no data near seed time", ha="center", va="center", fontsize=9, color="#7f8c8d")
            ax.set_title(f"cyl{config.CYLINDER_TO_INDEX[name]}", fontsize=9)
            ax.set_xticks([])
            ax.set_yticks([])
            continue
        seed_signal = seed_samples.get(int(row["id"]))
        if seed_signal is None or seed_signal.size == 0:
            ax.text(0.5, 0.5, "no sample data", ha="center", va="center", fontsize=9, color="#7f8c8d")
            continue
        fs = float(row["sample_frequency_hz"] or 0) or 25600.0
        base_signal = None
        if name in baseline_rows:
            base_signal = baseline_samples.get(int(baseline_rows[name]["id"]))
        scale = float(np.sqrt(np.mean(base_signal ** 2))) if base_signal is not None and base_signal.size else float(np.sqrt(np.mean(seed_signal ** 2)))
        scale = max(scale, 1e-9)
        x = np.arange(seed_signal.size) / fs * 1000.0
        if base_signal is not None and base_signal.size:
            xb = np.arange(base_signal.size) / fs * 1000.0
            ax.plot(xb, base_signal / scale, color="#95a5a6", linewidth=0.8, label="baseline")
        ax.plot(x, seed_signal / scale, color="#c0392b" if name == cylinder else "#2980b9", linewidth=0.9, label="seed")
        ax.set_title(f"cyl{config.CYLINDER_TO_INDEX[name]}", fontsize=9)
        ax.grid(alpha=0.2)
        ax.tick_params(labelsize=7)
        if name == cylinder:
            ax.set_ylabel("normalized amplitude", fontsize=8)
            ax.legend(loc="upper right", fontsize=7)
    fig.suptitle(f"{event['event_id']} cylinder#{config.CYLINDER_TO_INDEX[cylinder]} - first-cycle waveform (baseline vs seed)")
    fig.tight_layout()
    path = config.PLOT_DIR / f"waveforms_{event['event_id']}.png"
    fig.savefig(path, dpi=130)
    plt.close(fig)
    print(f"[plot] {path}", flush=True)