Compressor/backend/indicator_features/extract_from_db.py

#!/usr/bin/env python3
"""Extract per-cycle indicator-diagram features from the wave database (read-only).

Fetches a *small, sampled* set of running cycles (wave_sample_one holds one first
cycle per wave_file) for the chosen cylinder points and time range, computes the
phase-level features in indicator_features.py, and writes one CSV row per cycle.

Sampling keeps the query volume low: at most --per-day files per point per day,
picked evenly across that day's running files.
"""
from __future__ import annotations

import argparse
import csv
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path

import numpy as np
import pymysql

BACKEND = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND))

from app.config import settings  # noqa: E402
from indicator_features.indicator_features import (  # noqa: E402
    CYLINDER_STAGE,
    cycle_features,
)

CYLINDERS = ("一", "二", "三", "四", "五", "六")
SIDES = ("盖侧", "轴侧")


def connect() -> pymysql.connections.Connection:
    return pymysql.connect(
        host=settings.db_host, port=settings.db_port, user=settings.db_user,
        password=settings.db_password, database=settings.db_name, charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor, connect_timeout=8, read_timeout=300,
    )


def point_name(unit: str, cylinder: str, side: str) -> str:
    return f"{unit}号机组{cylinder}缸压力{side}"


def fetch_file_rows(connection, point: str, start: str, end: str) -> list[dict]:
    with connection.cursor() as cursor:
        cursor.execute(
            """
            SELECT id, sample_time, rpm
            FROM wave_file
            WHERE point_name = %s AND measurement_type = '压力'
              AND sample_time >= %s AND sample_time < %s
            ORDER BY sample_time
            """,
            (point, start, end),
        )
        return list(cursor.fetchall())


def sample_per_day(rows: list[dict], per_day: int) -> list[dict]:
    if per_day <= 0:
        return rows
    by_day: dict = defaultdict(list)
    for row in rows:
        by_day[row["sample_time"].date()].append(row)
    picked: list[dict] = []
    for day in sorted(by_day):
        day_rows = by_day[day]
        if len(day_rows) <= per_day:
            picked.extend(day_rows)
            continue
        idx = np.linspace(0, len(day_rows) - 1, per_day).round().astype(int)
        picked.extend(day_rows[int(i)] for i in idx)
    return picked


def fetch_cycles(connection, file_ids: list[int], batch: int) -> dict[int, np.ndarray]:
    cycles: dict[int, list[tuple[int, float]]] = defaultdict(list)
    with connection.cursor() as cursor:
        for start in range(0, len(file_ids), batch):
            chunk = file_ids[start:start + batch]
            placeholders = ", ".join(["%s"] * len(chunk))
            cursor.execute(
                f"""
                SELECT wave_file_id, sample_index, CAST(signal_value AS FLOAT) AS sig
                FROM wave_sample_one
                WHERE wave_file_id IN ({placeholders})
                ORDER BY wave_file_id, sample_index
                """,
                tuple(chunk),
            )
            for row in cursor.fetchall():
                if row["sig"] is None:
                    continue
                cycles[int(row["wave_file_id"])].append((int(row["sample_index"]), float(row["sig"])))
    result = {}
    for file_id, rows in cycles.items():
        rows.sort()
        result[file_id] = np.asarray([value for _, value in rows], dtype=float)
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description="只读抽取功图单周期特征（抽样）")
    parser.add_argument("--unit", required=True, choices=("7", "8", "9"))
    parser.add_argument("--cylinders", nargs="+", default=list(CYLINDERS), choices=CYLINDERS)
    parser.add_argument("--sides", nargs="+", default=list(SIDES), choices=SIDES)
    parser.add_argument("--start", required=True, help="YYYY-MM-DD")
    parser.add_argument("--end", required=True, help="YYYY-MM-DD（不含）")
    parser.add_argument("--per-day", type=int, default=3, help="每测点每天最多取几个文件")
    parser.add_argument("--batch", type=int, default=200, help="IN 查询批大小")
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    connection = connect()
    rows_out: list[dict] = []
    try:
        for cylinder in args.cylinders:
            for side in args.sides:
                point = point_name(args.unit, cylinder, side)
                files = fetch_file_rows(connection, point, args.start, args.end)
                picked = sample_per_day(files, args.per_day)
                print(f"{point}: 运行文件 {len(files)}，抽样 {len(picked)}", flush=True)
                if not picked:
                    continue
                id_to_time = {int(r["id"]): r for r in picked}
                cycles = fetch_cycles(connection, list(id_to_time), args.batch)
                for file_id, signal in cycles.items():
                    feats = cycle_features(signal, cylinder)
                    if not feats:
                        continue
                    meta = id_to_time[file_id]
                    rows_out.append({
                        "unit": args.unit,
                        "cylinder": cylinder,
                        "stage": CYLINDER_STAGE[cylinder],
                        "side": side,
                        "point_name": point,
                        "wave_file_id": file_id,
                        "sample_time": str(meta["sample_time"]),
                        "rpm": float(meta["rpm"] or 0.0),
                        **feats,
                    })
    finally:
        connection.close()

    if not rows_out:
        print("没有特征行，退出")
        return
    rows_out.sort(key=lambda r: (r["cylinder"], r["side"], r["sample_time"]))
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows_out[0].keys()))
        writer.writeheader()
        writer.writerows(rows_out)
    print(f"特征行 {len(rows_out):,} -> {output}")
    print(f"时间：{datetime.now():%Y-%m-%d %H:%M:%S}")


if __name__ == "__main__":
    main()