flow_control/data_logger.py

"""逐周期保存控制数据，并在运行结束时生成三联曲线图。"""

import csv
from datetime import datetime
import math
from pathlib import Path


CSV_FIELDS = (
    "time_s",
    "timestamp",
    "target_flow_slm",
    "measured_flow_slm",
    "opening_pct",
    "pressure_kpa",
    "error_slm",
    "motor_position",
    "actual_dt_s",
    "status",
)


class FlowRunRecorder:
    """把每个控制周期写入 CSV，并在结束时输出一张三联图。"""

    def __init__(self, output_directory, *, plot_dpi=160):
        output_dir = Path(output_directory)
        output_dir.mkdir(parents=True, exist_ok=True)

        run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.csv_path = output_dir / f"flow_run_{run_id}.csv"
        self.image_path = output_dir / f"flow_run_{run_id}.png"
        self.plot_dpi = int(plot_dpi)

        self._file = self.csv_path.open("w", newline="", encoding="utf-8-sig")
        self._writer = csv.DictWriter(self._file, fieldnames=CSV_FIELDS)
        self._writer.writeheader()
        self._file.flush()
        self._first_timestamp = None
        self.sample_count = 0
        self._closed = False

    def record(self, result):
        """保存一个 ControlStepResult；每次写入后立即刷新到磁盘。"""
        if self._closed:
            raise RuntimeError("记录器已经关闭")

        if self._first_timestamp is None:
            self._first_timestamp = float(result.timestamp)
        elapsed_s = max(0.0, float(result.timestamp) - self._first_timestamp)

        self._writer.writerow(
            {
                "time_s": f"{elapsed_s:.6f}",
                "timestamp": f"{float(result.timestamp):.6f}",
                "target_flow_slm": self._format_value(result.target_flow_slm),
                "measured_flow_slm": self._format_value(
                    result.measured_flow_slm
                ),
                "opening_pct": self._format_value(result.opening_pct),
                "pressure_kpa": self._format_value(result.pressure_kpa),
                "error_slm": self._format_value(result.error_slm),
                "motor_position": self._format_value(result.motor_position),
                "actual_dt_s": self._format_value(result.actual_dt_s),
                "status": result.status,
            }
        )
        self._file.flush()
        self.sample_count += 1

    def finalize(self):
        """关闭 CSV 并生成 PNG；无采样点时只保留 CSV。"""
        self.close()
        if self.sample_count == 0:
            return None
        self._create_plot()
        return self.image_path

    def close(self):
        if not self._closed:
            self._file.flush()
            self._file.close()
            self._closed = True

    def _create_plot(self):
        # 使用无界面后端，保证从终端运行和无显示器环境都能保存图片。
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        time_s = []
        target_flow = []
        measured_flow = []
        opening = []
        pressure = []

        with self.csv_path.open("r", newline="", encoding="utf-8-sig") as file:
            for row in csv.DictReader(file):
                time_s.append(float(row["time_s"]))
                target_flow.append(self._float_or_nan(row["target_flow_slm"]))
                measured_flow.append(
                    self._float_or_nan(row["measured_flow_slm"])
                )
                opening.append(self._float_or_nan(row["opening_pct"]))
                pressure.append(self._float_or_nan(row["pressure_kpa"]))

        figure, axes = plt.subplots(
            3,
            1,
            figsize=(12, 10),
            sharex=True,
            constrained_layout=True,
        )
        figure.suptitle("Gas Flow Control Response", fontsize=15)

        axes[0].plot(
            time_s,
            measured_flow,
            color="#1565C0",
            linewidth=1.4,
            label="Measured flow",
        )
        axes[0].plot(
            time_s,
            target_flow,
            color="#D84315",
            linewidth=1.5,
            linestyle="--",
            label="Target flow",
        )
        axes[0].set_ylabel("Flow (SLM)")
        axes[0].set_title("Flow response")
        axes[0].legend(loc="best")

        axes[1].plot(
            time_s,
            opening,
            color="#2E7D32",
            linewidth=1.4,
        )
        axes[1].set_ylabel("Opening (%)")
        axes[1].set_title("Valve opening")
        axes[1].set_ylim(-2, 102)

        if any(math.isfinite(value) for value in pressure):
            axes[2].plot(
                time_s,
                pressure,
                color="#6A1B9A",
                linewidth=1.4,
            )
        else:
            axes[2].text(
                0.5,
                0.5,
                "No pressure data",
                transform=axes[2].transAxes,
                ha="center",
                va="center",
                color="gray",
            )
        axes[2].set_ylabel("Pressure (kPa)")
        axes[2].set_xlabel("Time (s)")
        axes[2].set_title("Pressure")

        for axis in axes:
            axis.grid(True, alpha=0.3, linestyle="--")
            axis.margins(x=0)

        figure.savefig(self.image_path, dpi=self.plot_dpi, bbox_inches="tight")
        plt.close(figure)

    @staticmethod
    def _format_value(value):
        if value is None:
            return ""
        number = float(value)
        return "" if not math.isfinite(number) else f"{number:.6f}"

    @staticmethod
    def _float_or_nan(value):
        if value in (None, ""):
            return math.nan
        try:
            return float(value)
        except (TypeError, ValueError):
            return math.nan