ParkingRobot/data_process/旧版控制器轨迹测试处理/plot_speed_response.py

"""绘制控制器参考速度与Detour差分实际速度对比图。"""

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,
    segmented_savgol,
    shade_localization_jump_windows,
)


def calculate_actual_speed_mps(
    frame,
    filter_window_seconds: float,
) -> np.ndarray:
    """使用Savitzky-Golay求位置导数并计算Detour实际合速度。"""
    time = frame["TimeSeconds"].to_numpy(dtype=float)
    dt = float(np.median(np.diff(time)))
    # 直接对固定频率重采样后的位置做SG求导，避免“先平滑再求导”
    # 造成两次滤波和过度削弱速度峰值。
    x_mm = frame["DetourXRawMm"].to_numpy(dtype=float)
    y_mm = frame["DetourYRawMm"].to_numpy(dtype=float)
    vx_mm_per_second = segmented_savgol(
        x_mm,
        dt,
        filter_window_seconds,
        derivative=1,
    )
    vy_mm_per_second = segmented_savgol(
        y_mm,
        dt,
        filter_window_seconds,
        derivative=1,
    )

    speed = np.hypot(
        vx_mm_per_second,
        vy_mm_per_second,
    ) / 1000.0
    speed[
        frame["InvalidNearLocalizationJump"].to_numpy(dtype=bool)
    ] = np.nan
    return speed


def plot_speed(
    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)
    command_speed = frame["CommandSpeedMps"].to_numpy(dtype=float)
    actual_speed = calculate_actual_speed_mps(
        frame,
        filter_window_seconds,
    )
    is_in_place_rotation = (
        str(metadata["trajectory_name"])
        .lower()
        .startswith("rotate")
    )
    # 原地自转CSV中的ReferenceSpeed历史上保存的是角速度上限deg/s，
    # 不能作为线速度m/s使用；其参考线速度应为0。
    configured_speed = (
        0.0
        if is_in_place_rotation
        else float(metadata["reference_speed_mps"])
    )

    moving = (
        (command_speed > max(0.02, configured_speed * 0.1)) &
        np.isfinite(actual_speed)
    )
    if np.any(moving):
        speed_rmse = float(
            np.sqrt(
                np.mean(
                    (actual_speed[moving] - command_speed[moving]) ** 2
                )
            )
        )
    else:
        speed_rmse = float("nan")

    fig, ax = plt.subplots(figsize=(10.0, 5.8))
    ax.plot(
        time,
        command_speed,
        linewidth=1.8,
        label="控制器参考/下发线速度",
    )
    ax.plot(
        time,
        actual_speed,
        linewidth=1.5,
        label="Detour差分实际线速度（SG求导）",
    )
    ax.axhline(
        configured_speed,
        linestyle=":",
        linewidth=1.3,
        color="tab:green",
        label=(
            "原地自转参考线速度 0 m/s"
            if is_in_place_rotation
            else f"配置巡航速度 {configured_speed:.3f} m/s"
        ),
    )
    shade_localization_jump_windows(ax, metadata)
    ax.set_xlabel("时间 / s")
    ax.set_ylabel("线速度 / (m/s)")
    ax.set_title(
        f"参考速度与实际速度对比\n"
        f"{metadata['controller_name']} - "
        f"{metadata['trajectory_name']}，"
        f"运动段RMSE={speed_rmse:.4f} m/s"
    )
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()

    destination = output_path(
        csv_path,
        output_directory,
        "speed_response",
    )
    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="绘制参考速度与Detour差分实际速度对比图。"
    )
    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_speed(
            csv_path,
            args.frequency,
            args.window,
            args.output_dir,
            args.show,
        )
        print(f"已生成：{destination}")


if __name__ == "__main__":
    main()