Compressor/_tmp_cylinder_vibration/local_labels.py

"""动态局部基线：标签前最多 14 天、至少 1 天有效运行数据。"""
from __future__ import annotations

import numpy as np
import pandas as pd

from . import config, labels

MIN_BASELINE_POINTS = 96
WINDOW = "30D"
GAP_HOURS = 2
TRANSITION_HOURS = 1


def _transition_mask(group: pd.DataFrame) -> pd.Series:
    times = group["sample_time"].reset_index(drop=True)
    gaps = times.diff().dt.total_seconds() / 3600.0
    starts = gaps > GAP_HOURS
    mask = starts.copy()
    for offset in range(1, int(TRANSITION_HOURS * 4) + 1):
        mask |= starts.shift(offset, fill_value=False)
    return mask


def _local_baseline(frame: pd.DataFrame) -> pd.DataFrame:
    columns = ["rms_log", "kurt", "hf_ratio", "env_1x_log", "env_2x_log",
               "resid_rms_log", "resid_kurt", "resid_hf_ratio", "resid_env_1x_log", "resid_env_2x_log"]
    for column in columns:
        frame[f"base_{column}"] = np.nan
        frame[f"scale_{column}"] = np.nan
    frame["transition"] = False

    for _, group in frame.groupby("cylinder", sort=False):
        group = group.sort_values("sample_time")
        original_index = group.index
        mask = _transition_mask(group)
        frame.loc[original_index, "transition"] = mask.to_numpy()
        series_group = group.set_index("sample_time")
        for column in columns:
            series = series_group[column].where(~mask.to_numpy())
            median = series.rolling(WINDOW, min_periods=MIN_BASELINE_POINTS).median().shift(1)
            q1 = series.rolling(WINDOW, min_periods=MIN_BASELINE_POINTS).quantile(0.25).shift(1)
            q3 = series.rolling(WINDOW, min_periods=MIN_BASELINE_POINTS).quantile(0.75).shift(1)
            scale = (q3 - q1) / 1.349
            fallback = series.rolling(WINDOW, min_periods=MIN_BASELINE_POINTS).std().shift(1)
            scale = scale.where(scale > 1e-9, fallback)
            frame.loc[original_index, f"base_{column}"] = median.to_numpy()
            frame.loc[original_index, f"scale_{column}"] = scale.to_numpy()
    return frame


def build_local_labels(frame: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, dict[str, float]]:
    frame = labels._prepare_raw(frame)
    frame = _local_baseline(frame)

    score_columns = [f"resid_{name}" for name in config.RESIDUAL_FEATURES]
    for column in score_columns:
        frame[f"z_{column}"] = (frame[column] - frame[f"base_{column}"]) / frame[f"scale_{column}"]
        frame[f"z_{column}"] = frame[f"z_{column}"].replace([np.inf, -np.inf], np.nan).clip(-20, 20)

    z_columns = [f"z_{column}" for column in score_columns]
    z_values = frame[z_columns].abs()
    frame["score"] = z_values.max(axis=1).fillna(0.0)
    frame["driver"] = z_values.fillna(0.0).idxmax(axis=1).str.replace("z_", "", regex=False)
    frame.loc[frame["score"] <= 0, "driver"] = "none"
    baseline_columns = [f"base_{column}" for column in score_columns] + [f"scale_{column}" for column in score_columns]
    frame["baseline_ok"] = frame[baseline_columns].notna().all(axis=1)

    frame["rms_ratio"] = np.power(10.0, frame["rms_log"] - frame["base_rms_log"])
    frame["kurt_ratio"] = frame["kurt"] / frame["base_kurt"]

    frame["contrast"] = 0.0
    for _, group in frame.groupby("cylinder", sort=False):
        group = group.sort_values("sample_time")
        score = group.set_index("sample_time")["score"]
        local = score.rolling("24h", min_periods=4).median().shift(1)
        frame.loc[group.index, "contrast"] = (score - local).fillna(0.0).to_numpy()

    thresholds: dict[str, float] = {}
    anomaly_thresholds: dict[str, float] = {}
    severe_thresholds: dict[str, float] = {}
    for cylinder, group in frame.groupby("cylinder", sort=False):
        normal = group[(group["baseline_ok"]) & (~group["transition"])]["score"].dropna()
        if normal.empty:
            normal = group["score"].dropna()
        median = float(normal.median()) if not normal.empty else 1.0
        mad = float((normal - median).abs().median()) if not normal.empty else 0.0
        mad = max(mad, 0.25)
        candidate = max(3.5, median + 6.0 * mad)
        anomaly = max(5.0, median + 10.0 * mad)
        severe = max(7.0, median + 14.0 * mad)
        thresholds[cylinder] = candidate
        anomaly_thresholds[cylinder] = max(candidate, anomaly)
        severe_thresholds[cylinder] = max(candidate, anomaly, severe)
    frame["threshold"] = frame["cylinder"].map(thresholds)
    frame["anomaly_threshold"] = frame["cylinder"].map(anomaly_thresholds)
    frame["severe_threshold"] = frame["cylinder"].map(severe_thresholds)
    frame["hit"] = frame["baseline_ok"] & (frame["score"] >= frame["threshold"])

    point_labels = labels._extract_segments(frame, 15)
    return frame, point_labels, thresholds