Compressor/_tmp_cylinder_vibration/features.py

"""首周期波形特征。"""
from __future__ import annotations

import numpy as np
from scipy.signal import hilbert

EPS = 1e-12
BANDS = [
    ("band_0_1", 0.5, 1.5),
    ("band_1_2", 1.5, 2.5),
    ("band_2_5", 2.5, 5.0),
    ("band_5_10", 5.0, 10.0),
    ("band_10_20", 10.0, 20.0),
    ("band_20_40", 20.0, 40.0),
    ("band_40_80", 40.0, 80.0),
]


def _safe_log(value: float) -> float:
    return float(np.log10(max(value, 0.0) + EPS))


def _band_power(freqs: np.ndarray, power: np.ndarray, low: float, high: float) -> float:
    mask = (freqs >= low) & (freqs < high)
    return float(power[mask].sum())


def _band_amplitude(freqs: np.ndarray, amplitude: np.ndarray, center: float) -> float:
    if center <= 0:
        return 0.0
    low, high = center * 0.8, center * 1.2
    mask = (freqs >= low) & (freqs <= high)
    if not mask.any():
        return 0.0
    return float(amplitude[mask].max())


def extract_features(signal: np.ndarray, fs: float, rpm: float) -> dict | None:
    x = np.asarray(signal, dtype=float)
    x = x[np.isfinite(x)]
    if x.size < 32:
        return None

    fs = float(fs or 0.0)
    rpm = float(rpm or 0.0)
    if fs <= 0:
        return None

    x = x - x.mean()
    n = x.size
    std = float(x.std())
    rms = float(np.sqrt(np.mean(x * x)))
    peak = float(np.max(np.abs(x)))
    p2p = float(x.max() - x.min())
    mean_abs = float(np.mean(np.abs(x)))
    kurt = float(np.mean((x / std) ** 4)) if std > EPS else 0.0
    skew = float(np.mean((x / std) ** 3)) if std > EPS else 0.0
    crest = peak / rms if rms > EPS else 0.0
    impulse = peak / mean_abs if mean_abs > EPS else 0.0
    clearance = peak / (mean_abs ** 2) if mean_abs > EPS else 0.0
    shape = rms / mean_abs if mean_abs > EPS else 0.0
    zcr = float(np.mean(np.diff(np.signbit(x).astype(np.int8)) != 0))

    window = np.hanning(n)
    spectrum = np.abs(np.fft.rfft(x * window)) / max(n, 1)
    power = spectrum ** 2
    freqs = np.fft.rfftfreq(n, d=1.0 / fs)
    total_power = float(power[freqs > 0].sum())
    if total_power > EPS:
        probabilities = power[freqs > 0] / total_power
        spec_entropy = float(-np.sum(probabilities * np.log(probabilities + EPS)) / np.log(len(probabilities)))
        spec_centroid = float(np.sum(freqs[freqs > 0] * probabilities) / 1000.0)
    else:
        spec_entropy = 0.0
        spec_centroid = 0.0

    f0 = rpm / 60.0
    band_features: dict[str, float] = {}
    for name, low, high in BANDS:
        if f0 > 0:
            energy = _band_power(freqs, power, low * f0, high * f0)
        else:
            energy = 0.0
        band_features[f"{name}_log"] = _safe_log(energy)

    if total_power > EPS:
        threshold = 10.0 * f0 if f0 > 0 else 2000.0
        hf_ratio = float(power[freqs >= threshold].sum() / total_power)
    else:
        hf_ratio = 0.0

    envelope = np.abs(hilbert(x))
    env_centered = envelope - envelope.mean()
    env_window = np.hanning(n)
    env_spectrum = np.abs(np.fft.rfft(env_centered * env_window)) / max(n, 1)
    env_power = env_spectrum ** 2
    env_total = float(env_power[freqs > 0].sum())
    if env_total > EPS:
        env_probabilities = env_power[freqs > 0] / env_total
        env_entropy = float(-np.sum(env_probabilities * np.log(env_probabilities + EPS)) / np.log(len(env_probabilities)))
    else:
        env_entropy = 0.0

    env_1x = _band_amplitude(freqs, env_spectrum, f0)
    env_2x = _band_amplitude(freqs, env_spectrum, 2.0 * f0)
    env_3x = _band_amplitude(freqs, env_spectrum, 3.0 * f0)

    return {
        "n_samples": int(n),
        "rms": rms,
        "rms_log": _safe_log(rms),
        "peak": peak,
        "p2p": p2p,
        "std": std,
        "skew": skew,
        "kurt": kurt,
        "crest": crest,
        "impulse": impulse,
        "clearance": clearance,
        "shape": shape,
        "zcr": zcr,
        "spec_centroid": spec_centroid,
        "spec_entropy": spec_entropy,
        "hf_ratio": hf_ratio,
        "env_rms_log": _safe_log(float(np.sqrt(np.mean(envelope ** 2)))),
        "env_peak": float(envelope.max()),
        "env_1x_log": _safe_log(env_1x),
        "env_2x_log": _safe_log(env_2x),
        "env_3x_log": _safe_log(env_3x),
        "env_entropy": env_entropy,
        **band_features,
    }