ParkingRobot/data_process/plot_angular_command.py

"""绘制控制器下发角速度命令曲线。"""

from __future__ import annotations

import argparse
from pathlib import Path

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

from plot_trajectory_comparison import (
    configure_matplotlib,
    discover_csv_files,
    load_and_resample,
    output_path,
    shade_localization_jump_windows,
)


def plot_angular_command(
    csv_path: Path,
    frequency_hz: float,
    filter_window_seconds: float,
    output_directory: str | None,
    show: bool,
) -> Path:
    """生成单份CSV的命令角速度曲线。"""
    frame, metadata = load_and_resample(
        csv_path,
        frequency_hz,
        filter_window_seconds,
    )
    time = frame["TimeSeconds"].to_numpy(dtype=float)
    angular_command = frame[
        "CommandAngularSpeedRadPerSec"
    ].to_numpy(dtype=float)
    maximum = float(np.max(angular_command))
    minimum = float(np.min(angular_command))

    fig, ax = plt.subplots(figsize=(10.0, 5.5))
    ax.plot(
        time,
        angular_command,
        color="tab:red",
        linewidth=1.6,
        label="CommandAngularSpeed",
    )
    ax.axhline(0.0, color="black", linewidth=0.8)
    shade_localization_jump_windows(ax, metadata)
    ax.set_xlabel("时间 / s")
    ax.set_ylabel("命令角速度 / (rad/s)")
    ax.set_title(
        f"角速度指令曲线\n"
        f"{metadata['controller_name']} - "
        f"{metadata['trajectory_name']}，"
        f"范围=[{minimum:.3f}, {maximum:.3f}]rad/s"
    )
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()

    destination = output_path(
        csv_path,
        output_directory,
        "angular_command",
    )
    fig.savefig(destination, dpi=300, bbox_inches="tight")
    if show:
        plt.show()
    plt.close(fig)
    return destination


def main() -> None:
    configure_matplotlib()
    parser = argparse.ArgumentParser(
        description="绘制控制器下发角速度命令曲线。"
    )
    parser.add_argument("files", nargs="*", help="一个或多个CSV文件")
    parser.add_argument("--frequency", type=float, default=20.0)
    parser.add_argument("--window", type=float, default=0.55)
    parser.add_argument("--output-dir")
    parser.add_argument("--show", action="store_true")
    args = parser.parse_args()

    for csv_path in discover_csv_files(args.files):
        destination = plot_angular_command(
            csv_path,
            args.frequency,
            args.window,
            args.output_dir,
            args.show,
        )
        print(f"已生成：{destination}")


if __name__ == "__main__":
    main()