aggregate.py

import sys
from argparse import ArgumentParser
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

# Set Times New Roman as the default font
plt.rcParams["font.family"] = "Times New Roman"
plt.rcParams["mathtext.fontset"] = "custom"
plt.rcParams["mathtext.rm"] = "Times New Roman"
plt.rcParams["mathtext.it"] = "Times New Roman:italic"
plt.rcParams["mathtext.bf"] = "Times New Roman:bold"


def load(file_path: Path) -> np.ndarray:
    data = []
    with open(file_path, "r") as input:
        for line in input:
            line = line.strip()
            if len(line) == 0 or line.startswith("#"):
                continue
            try:
                data.append([float(x) for x in line.split(",")])
            except ValueError:
                continue
    data.sort(key=lambda x: x[0])
    if len(data) == 0:
        raise ValueError("No valid trajectory points found.")
    arr = np.array(data, dtype=np.float64)
    arr[:, 0] -= arr[0, 0]  # Normalize time to start from 0
    return arr


parser = ArgumentParser(description="Configuration for the application")
parser.add_argument(
    "folders",
    nargs="+",
    help="Input coverage records",
)
parser.add_argument(
    "--time-resolution",
    type=float,
    default=1.0,
    help="Time resolution in seconds, default is 1.0s",
)
parser.add_argument(
    "--time-limit",
    type=float,
    default=None,
    help="Time limit in seconds, default is None (no limit)",
)
parser.add_argument(
    "--display",
    action="store_true",
    help="Display plot of results",
)
parser.add_argument(
    "--export",
    type=Path,
    default=None,
    help="Export aggregated plot to specified path",
)


args = parser.parse_args()
FOLDERS: list[str] = list(args.folders)
TIME_RES = float(args.time_resolution)
DISPLAY: bool = bool(args.display)
EXPORT_PATH: Path | None = args.export
T_LIM: float | None = args.time_limit

AGGREGATED: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = {}


def aggregate(dir: Path):
    if not dir.is_dir():
        raise ValueError(f"Input path {dir} is not a directory.")
    INPUT_FILES: list[Path] = list(dir.glob("**/coverage.list"))
    print(f"Found {len(INPUT_FILES)} coverage records in {dir}")
    if len(INPUT_FILES) == 0:
        raise ValueError(f"No coverage records found in {dir}.")
    DATASET: list[np.ndarray] = [load(fp) for fp in INPUT_FILES]

    T1 = max(arr[-1, 0] for arr in DATASET)  # seconds

    T = np.arange(0.0, T1, TIME_RES)
    BATCHES: list[list[float]] = []

    for ddl in T:
        # Collect samples from each trajectory within [0, ddl)
        # Use the last sample before ddl
        batch = []
        for arr in DATASET:
            idx = np.searchsorted(arr[:, 0], ddl, side="right") - 1
            if idx >= 0:
                batch.append(arr[idx, 1])
        BATCHES.append(batch)

    # Compute mean and stddev for each batch
    MEAN = [np.mean(B) for B in BATCHES]
    STDDEV = [np.std(B) for B in BATCHES]

    AGGREGATED[dir.name] = (np.array(T), np.array(MEAN), np.array(STDDEV))

    with open(dir / "coverage.aggregated.list", "w") as output:
        print("# time,mean,stddev", file=output)
        for t, mean, std in zip(T, MEAN, STDDEV):
            print(t, mean, std, sep=",", file=output)

    # Plot
    fig, (ax1, ax2) = plt.subplots(2, 1)
    if fig.canvas.manager is not None:
        fig.canvas.manager.set_window_title(dir.name)
    else:
        fig.suptitle(dir.name)

    # Plot trajectories and store line objects with their file paths
    lines_with_paths = []
    for arr, file_path in zip(DATASET, INPUT_FILES):
        (line,) = ax1.plot(arr[:, 0], arr[:, 1], linestyle="-", label="Trajectory")
        lines_with_paths.append((line, file_path))

    ax1.set_title("Trajectories")
    ax1.set_xlabel("Time (s)")
    ax1.set_ylabel("Coverage ($m^2$)")
    ax1.grid(True)

    # Add hover annotation
    annot = ax1.annotate(
        "",
        xy=(0, 0),
        xytext=(10, 10),
        textcoords="offset points",
        bbox=dict(boxstyle="round", fc="yellow", alpha=0.8),
        arrowprops=dict(arrowstyle="->"),
    )
    annot.set_visible(False)

    def on_hover(event):
        if event.inaxes == ax1:
            for line, file_path in lines_with_paths:
                cont, _ = line.contains(event)
                if cont:
                    annot.xy = (event.xdata, event.ydata)
                    annot.set_text(str(file_path))
                    annot.set_visible(True)
                    fig.canvas.draw_idle()
                    return
        annot.set_visible(False)
        fig.canvas.draw_idle()

    fig.canvas.mpl_connect("motion_notify_event", on_hover)

    ax2.plot(T, MEAN, "-", color="blue", label="Mean")
    ax2.fill_between(
        T,
        np.array(MEAN) - np.array(STDDEV),
        np.array(MEAN) + np.array(STDDEV),
        alpha=0.3,
        color="blue",
        label="±1 Stddev",
    )
    ax2.set_title("Mean and Stddev over Time")
    ax2.set_xlabel("Time (s)")
    ax2.set_ylabel("Coverage ($m^2$)")
    ax2.legend()
    ax2.grid(True)
    if DISPLAY:
        plt.show(block=False)

    fig.savefig(dir / "coverage.aggregated.pdf")


for folder in FOLDERS:
    aggregate(Path(folder))


def desc(name: str):
    if name.endswith(".NO_FAM"):
        return "Familiarity Disabled"
    elif name.endswith(".NO_LKA"):
        return "Look Around Disabled"
    else:
        return "Full System"


# Find out max duration T1 among all aggregated results
if T_LIM is None:
    T_MAX = 0.0
    for name, (T, MEAN, STDDEV) in AGGREGATED.items():
        if T[-1] > T_MAX:
            T_MAX = T[-1]
else:
    T_MAX = T_LIM

if not DISPLAY and EXPORT_PATH is None:
    sys.exit(0)

# Plot all aggregated results together
fig, ax = plt.subplots(figsize=(8, 4))
if fig.canvas.manager is not None:
    fig.canvas.manager.set_window_title("Aggregated Coverage")
OFFSET, Y_MAX = (
    min(min(Y) for (_, Y, _) in AGGREGATED.values()),
    max(max(Y + S) for (_, Y, S) in AGGREGATED.values()),
)
for name, (T, MEAN, STDDEV) in AGGREGATED.items():
    T = np.append(T, T_MAX)
    MEAN = np.append(MEAN, MEAN[-1]) - OFFSET
    STDDEV = np.append(STDDEV, STDDEV[-1])
    ax.plot(T, MEAN, "-", label=f"{desc(name)}")
    ax.fill_between(
        T,
        MEAN - STDDEV,
        MEAN + STDDEV,
        alpha=0.2,
    )
ax.set_xlabel("Time ($s$)")
ax.set_ylabel("Coverage ($m^2$)")
ax.set_xlim(0, T_MAX)
ax.set_ylim(0, Y_MAX)
ax.legend()
ax.grid(True)
if EXPORT_PATH is not None:
    fig.savefig(EXPORT_PATH)

if DISPLAY:
    plt.show(block=True)