Compressor/models/generalized/build_datasets.py

#!/usr/bin/env python3
"""Build pressure model samples from the reconstructed pressure algorithm output."""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import paths, queries
from models.common import windowing
from models.generalized import features_pressure, pressure_algorithm
from models.common.constants import MIN_HOUR_COVERAGE


def cycle_baseline(series, days=3):
    return windowing.cycle_baseline(series, days)


def pressure_features(series, end_index, window):
    from models.common.constants import HOURLY_BASELINE_HOURS

    oil = [row["oil_avg"] for row in window]
    coverages = [row["运行覆盖率"] for row in window]
    baseline = windowing.median_value([row["oil_avg"] for row in window[:-1]])
    current = oil[-1]
    oil_mean, oil_min, oil_max = windowing.mean(oil), min(oil), max(oil)
    oil_range = windowing.value_range(oil)
    oil_slope = windowing.slope(oil)
    cycle_base = cycle_baseline(series)
    recent_center, _ = windowing.previous_baseline(series, end_index, "oil_avg", HOURLY_BASELINE_HOURS)
    return {
        "oil_mean": oil_mean, "oil_min": oil_min, "oil_max": oil_max, "oil_range": oil_range,
        "oil_slope": oil_slope,
        "oil_last_relative_to_window": (current / baseline - 1) if baseline else None,
        "oil_low_ratio_mean": windowing.mean([windowing.ratio(row["oil_low_count"], row["运行样本数"]) for row in window]),
        "oil_low_low_ratio_mean": windowing.mean([windowing.ratio(row["oil_low_low_count"], row["运行样本数"]) for row in window]),
        "running_coverage_mean": windowing.mean(coverages), "running_coverage_min": min(coverages),
        "oil_rel_cycle_mean": (oil_mean / cycle_base - 1) if cycle_base else None,
        "oil_rel_cycle_min": (oil_min / cycle_base - 1) if cycle_base else None,
        "oil_rel_cycle_last": (current / cycle_base - 1) if cycle_base else None,
        "oil_rel_recent_mean": (oil_mean / recent_center - 1) if recent_center else None,
        "oil_rel_recent_last": (current / recent_center - 1) if recent_center else None,
        "oil_slope_rel_cycle": (oil_slope / cycle_base) if cycle_base else None,
        "oil_decline_hours": windowing.decline_hours(series, end_index),
        "oil_drawdown_cycle": ((cycle_base - oil_min) / cycle_base) if cycle_base else None,
        "oil_range_rel_cycle": (oil_range / cycle_base) if cycle_base else None,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="构建压力模型样本（窗口可配置）")
    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("--window-hours", type=int, default=3)
    parser.add_argument("--output-dir", default=str(paths.PRESSURE_DATAS))
    parser.add_argument("--hourly-csv", default=str(paths.HOURLY_FULL))
    parser.add_argument("--units", nargs="+", default=None)
    args = parser.parse_args()

    output = Path(args.output_dir)
    if output.exists():
        shutil.rmtree(output)
    output.mkdir(parents=True)

    hourly_path = Path(args.hourly_csv)
    if hourly_path.is_file():
        print(f"复用小时中间表：{hourly_path}")
        rows = queries.load_hourly_csv(hourly_path)
    else:
        from models.common import db

        batch_to_unit = db.load_pks_batches()
        selected = sorted(batch_to_unit)
        configs = db.load_alarm_config({batch: batch_to_unit[batch] for batch in selected})
        start, end = queries.parse_dt(args.start), queries.parse_dt(args.end)
        rows = queries.fetch_all_hourly(start, end, args.chunk_days, batch_to_unit, configs, selected)

    scoped = [row for row in rows if not args.units or row["机组"] in set(args.units)]
    windowing.write_hourly(output / "hourly.csv", scoped)
    stages = windowing.stage_ranges(paths.PRESSURE_RESULTS / "润滑油压力小时阶段段.csv", "压力阶段")
    if args.units:
        stages = [stage for stage in stages if stage["unit"] in set(args.units)]
    samples = windowing.build_samples(scoped, "pressure", stages, args.window_hours, pressure_features)
    counts = windowing.write_samples(output, samples)
    print(f"pressure samples: {dict(counts)}")
    print(f"小时中间数据：{len(scoped):,} 行")
    print(f"样本输出目录：{output}")


if __name__ == "__main__":
    main()