ParkingRobot/data_process/plot_tracking_errors.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 (
    build_reference,
    configure_matplotlib,
    discover_csv_files,
    load_and_resample,
    output_path,
    shade_localization_jump_windows,
)


def plot_errors(
    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,
    )
    reference = build_reference(frame, metadata)
    time = frame["TimeSeconds"].to_numpy()
    lateral = np.asarray(reference["lateral_error_mm"])
    heading = np.asarray(reference["heading_error_degrees"])
    invalid = frame[
        "InvalidNearLocalizationJump"
    ].to_numpy(dtype=bool)
    lateral_for_statistics = lateral.copy()
    heading_for_statistics = heading.copy()
    lateral_for_statistics[invalid] = np.nan
    heading_for_statistics[invalid] = np.nan

    lateral_rmse = float(
        np.sqrt(np.nanmean(lateral_for_statistics**2))
    )
    heading_rmse = float(
        np.sqrt(np.nanmean(heading_for_statistics**2))
    )
    lateral_max = float(
        np.nanmax(np.abs(lateral_for_statistics))
    )
    heading_max = float(
        np.nanmax(np.abs(heading_for_statistics))
    )
    is_in_place_rotation = (
        reference["kind"] == "in_place_rotation"
    )

    fig, axes = plt.subplots(
        2,
        1,
        figsize=(10.0, 7.0),
        sharex=True,
    )
    axes[0].plot(time, lateral, linewidth=1.5)
    axes[0].axhline(0.0, color="black", linewidth=0.8)
    if is_in_place_rotation:
        axes[0].set_ylabel("旋转中心位置漂移 / mm")
        axes[0].set_title(
            f"原地自转位置漂移：RMS={lateral_rmse:.2f} mm，"
            f"最大值={lateral_max:.2f} mm"
        )
    else:
        axes[0].set_ylabel("横向误差 / mm")
        axes[0].set_title(
            f"横向误差：RMSE={lateral_rmse:.2f} mm，"
            f"最大绝对值={lateral_max:.2f} mm"
        )
    shade_localization_jump_windows(axes[0], metadata)
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(
        time,
        heading,
        color="tab:orange",
        linewidth=1.5,
    )
    axes[1].axhline(0.0, color="black", linewidth=0.8)
    axes[1].set_xlabel("时间 / s")
    axes[1].set_ylabel(
        "目标角度剩余误差 / °"
        if is_in_place_rotation
        else "航向误差 / °"
    )
    axes[1].set_title(
        (
            f"目标角度剩余误差：RMSE={heading_rmse:.2f}°，"
            f"最大绝对值={heading_max:.2f}°"
        )
        if is_in_place_rotation
        else (
            f"航向误差：RMSE={heading_rmse:.2f}°，"
            f"最大绝对值={heading_max:.2f}°"
        )
    )
    shade_localization_jump_windows(axes[1], metadata)
    axes[1].grid(True, alpha=0.3)
    if metadata["localization_jump_events"]:
        axes[1].legend(loc="best")

    fig.suptitle(
        f"横向/航向误差随时间变化\n"
        f"{metadata['controller_name']} - "
        f"{metadata['trajectory_name']}"
    )
    fig.tight_layout()

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

    if is_in_place_rotation:
        print(
            f"{csv_path.name}: position drift RMS="
            f"{lateral_rmse:.3f} mm, "
            f"target-angle error RMS={heading_rmse:.3f} deg"
        )
    else:
        print(
            f"{csv_path.name}: lateral RMSE="
            f"{lateral_rmse:.3f} mm, "
            f"heading RMSE={heading_rmse:.3f} deg"
        )
    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_errors(
            csv_path,
            args.frequency,
            args.window,
            args.output_dir,
            args.show,
        )
        print(f"已生成：{destination}")


if __name__ == "__main__":
    main()