spu stuff

SW/spu/scripts/plot_motor_ticks.py

@@ -0,0 +1,357 @@
+#!/usr/bin/env python3
+import argparse
+import struct
+import threading
+import time
+from collections import deque
+from queue import Empty, Queue
+
+try:
+    import pygame
+except ImportError as exc:
+    raise SystemExit(
+        "Missing dependency: pygame. Install with `pip install pygame pyserial`."
+    ) from exc
+
+try:
+    import serial
+except ImportError as exc:
+    raise SystemExit(
+        "Missing dependency: pyserial. Install with `pip install pygame pyserial`."
+    ) from exc
+
+
+RECORD_STRUCT = struct.Struct("<IIIII")
+RECORD_FIELDS = ("ms", "ticks_r", "target_r", "ticks_l", "target_l")
+
+BACKGROUND = (18, 20, 24)
+GRID = (48, 54, 64)
+AXIS = (140, 148, 160)
+TEXT = (230, 235, 240)
+ERROR = (255, 110, 110)
+
+TRACE_DEFS = [
+    {
+        "name": "speed_r",
+        "label": "SPEED_R",
+        "color": (80, 220, 140),
+        "unit": "ticks/s",
+        "kind": "derived",
+        "compute": lambda window: compute_speed(window, "ticks_r"),
+    },
+    {
+        "name": "target_r",
+        "label": "TARGET_R",
+        "color": (220, 80, 140),
+        "unit": "ticks/s",
+        "kind": "field",
+        "field": "target_r",
+    },
+    {
+        "name": "speed_l",
+        "label": "SPEED_L",
+        "color": (80, 170, 255),
+        "unit": "ticks/s",
+        "kind": "derived",
+        "compute": lambda window: compute_speed(window, "ticks_l"),
+    },
+    {
+        "name": "target_l",
+        "label": "TARGET_L",
+        "color": (255, 190, 80),
+        "unit": "ticks/s",
+        "kind": "field",
+        "field": "target_l",
+    },
+]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Read binary motor samples from a serial port and draw traces live."
+    )
+    parser.add_argument("port", help="Serial port, for example /dev/ttyUSB0 or COM3")
+    parser.add_argument(
+        "-b",
+        "--baudrate",
+        type=int,
+        default=115200,
+        help="Serial baudrate (default: 115200)",
+    )
+    parser.add_argument(
+        "--history-seconds",
+        type=float,
+        default=10.0,
+        help="Initial visible history window in seconds (default: 10)",
+    )
+    parser.add_argument(
+        "--avg-window",
+        type=int,
+        default=5,
+        help="Sliding average window in samples for derived traces (default: 5)",
+    )
+    parser.add_argument(
+        "--max-samples",
+        type=int,
+        default=30000,
+        help="Maximum number of samples to keep in memory (default: 30000)",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=1200,
+        help="Window width in pixels (default: 1200)",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=700,
+        help="Window height in pixels (default: 700)",
+    )
+    return parser.parse_args()
+
+
+def serial_reader(port: str, baudrate: int, output: Queue) -> None:
+    try:
+        with serial.Serial(port, baudrate=baudrate, timeout=1) as ser:
+            while True:
+                raw = ser.read(RECORD_STRUCT.size)
+                if len(raw) != RECORD_STRUCT.size:
+                    continue
+                values = RECORD_STRUCT.unpack(raw)
+                sample = dict(zip(RECORD_FIELDS, values))
+                sample["wall_time"] = time.monotonic()
+                output.put(sample)
+    except serial.SerialException as exc:
+        output.put({"error": f"Serial error: {exc}"})
+
+
+def compute_speed(
+    raw_window: list[dict[str, float | int]],
+    tick_field: str,
+) -> float | None:
+    if len(raw_window) < 2:
+        return None
+
+    first = raw_window[0]
+    last = raw_window[-1]
+    dt_ms = int(last["ms"]) - int(first["ms"])
+    if dt_ms <= 0:
+        return None
+
+    dt_ticks = int(last[tick_field]) - int(first[tick_field])
+    return (1000.0 * dt_ticks) / dt_ms
+
+
+def update_traces(
+    raw_samples: deque[dict[str, float | int]],
+    trace_samples: dict[str, deque[tuple[float, float, dict[str, float | int]]]],
+    avg_window: int,
+) -> None:
+    window = list(raw_samples)[-max(avg_window, 2):]
+    sample = raw_samples[-1]
+
+    for trace in TRACE_DEFS:
+        value = None
+        if trace["kind"] == "field":
+            value = float(sample[trace["field"]])
+        elif trace["kind"] == "derived":
+            value = trace["compute"](window)
+
+        if value is not None:
+            trace_samples[trace["name"]].append(
+                (float(sample["wall_time"]), float(value), sample)
+            )
+
+
+def drain_queue(
+    queue: Queue,
+    raw_samples: deque[dict[str, float | int]],
+    trace_samples: dict[str, deque[tuple[float, float, dict[str, float | int]]]],
+    avg_window: int,
+) -> str | None:
+    error_message = None
+    while True:
+        try:
+            item = queue.get_nowait()
+        except Empty:
+            return error_message
+
+        if "error" in item:
+            error_message = str(item["error"])
+            continue
+
+        raw_samples.append(item)
+        update_traces(raw_samples, trace_samples, avg_window)
+
+
+def draw_grid(
+    surface: pygame.Surface,
+    rect: pygame.Rect,
+    font: pygame.font.Font,
+    view_span: float,
+    y_max: float,
+) -> None:
+    pygame.draw.rect(surface, GRID, rect, width=1)
+
+    for fraction in (0.0, 0.25, 0.5, 0.75, 1.0):
+        y = rect.top + round(fraction * rect.height)
+        pygame.draw.line(surface, GRID, (rect.left, y), (rect.right, y), 1)
+        value = (1.0 - fraction) * y_max
+        text = font.render(f"{value:.1f}", True, AXIS)
+        surface.blit(text, (10, y - text.get_height() // 2))
+
+    for fraction in (0.0, 0.25, 0.5, 0.75, 1.0):
+        x = rect.left + round(fraction * rect.width)
+        pygame.draw.line(surface, GRID, (x, rect.top), (x, rect.bottom), 1)
+        seconds_ago = view_span * (1.0 - fraction)
+        label = f"-{seconds_ago:.1f}s" if seconds_ago > 0.05 else "now"
+        text = font.render(label, True, AXIS)
+        surface.blit(text, (x - text.get_width() // 2, rect.bottom + 8))
+
+
+def draw_trace(
+    surface: pygame.Surface,
+    rect: pygame.Rect,
+    samples: deque[tuple[float, float, dict[str, float | int]]],
+    view_span: float,
+    now: float,
+    color: tuple[int, int, int],
+    y_max: float,
+) -> tuple[float | None, dict[str, float | int] | None]:
+    visible_points = []
+    latest_value = None
+    latest_sample = None
+
+    for timestamp, value, sample in samples:
+        age = now - timestamp
+        if age < 0 or age > view_span:
+            continue
+        x = rect.right - (age / view_span) * rect.width
+        y = rect.bottom - (value / y_max) * rect.height
+        visible_points.append((round(x), round(y)))
+        latest_value = value
+        latest_sample = sample
+
+    if len(visible_points) >= 2:
+        pygame.draw.lines(surface, color, False, visible_points, 2)
+    elif len(visible_points) == 1:
+        pygame.draw.circle(surface, color, visible_points[0], 2)
+
+    return latest_value, latest_sample
+
+
+def zoom(view_span: float, direction: int, initial_span: float) -> float:
+    if direction > 0:
+        return max(1.0, view_span / 1.2)
+    return min(max(initial_span * 20.0, 60.0), view_span * 1.2)
+
+
+def main() -> int:
+    args = parse_args()
+    raw_samples: deque[dict[str, float | int]] = deque(maxlen=args.max_samples)
+    trace_samples = {
+        trace["name"]: deque(maxlen=args.max_samples) for trace in TRACE_DEFS
+    }
+    queue: Queue = Queue()
+
+    thread = threading.Thread(
+        target=serial_reader, args=(args.port, args.baudrate, queue), daemon=True
+    )
+    thread.start()
+
+    pygame.init()
+    pygame.display.set_caption(f"Motor traces: {args.port} @ {args.baudrate}")
+    screen = pygame.display.set_mode((args.width, args.height), pygame.RESIZABLE)
+    font = pygame.font.SysFont("monospace", 18)
+    small_font = pygame.font.SysFont("monospace", 14)
+    clock = pygame.time.Clock()
+
+    view_span = max(args.history_seconds, 1.0)
+    error_message = None
+    running = True
+
+    while running:
+        for event in pygame.event.get():
+            if event.type == pygame.QUIT:
+                running = False
+            elif event.type == pygame.MOUSEWHEEL:
+                view_span = zoom(view_span, event.y, args.history_seconds)
+            elif event.type == pygame.MOUSEBUTTONDOWN:
+                if event.button == 4:
+                    view_span = zoom(view_span, 1, args.history_seconds)
+                elif event.button == 5:
+                    view_span = zoom(view_span, -1, args.history_seconds)
+
+        queued_error = drain_queue(queue, raw_samples, trace_samples, args.avg_window)
+        if queued_error is not None:
+            error_message = queued_error
+
+        width, height = screen.get_size()
+        plot_rect = pygame.Rect(70, 40, max(100, width - 100), max(100, height - 100))
+        now = time.monotonic()
+
+        y_max = 1.0
+        for trace in TRACE_DEFS:
+            for timestamp, value, _sample in trace_samples[trace["name"]]:
+                age = now - timestamp
+                if 0 <= age <= view_span and value > y_max:
+                    y_max = value
+        y_max = max(10.0, y_max * 1.2)
+
+        screen.fill(BACKGROUND)
+        draw_grid(screen, plot_rect, small_font, view_span, y_max)
+
+        latest_values: dict[str, float] = {}
+        latest_sample: dict[str, float | int] | None = None
+
+        for trace in TRACE_DEFS:
+            value, sample = draw_trace(
+                screen,
+                plot_rect,
+                trace_samples[trace["name"]],
+                view_span,
+                now,
+                trace["color"],
+                y_max,
+            )
+            if value is not None:
+                latest_values[trace["name"]] = value
+            if sample is not None:
+                latest_sample = sample
+
+        header = (
+            f"port={args.port} baud={args.baudrate} zoom={view_span:.1f}s "
+            f"avg_window={args.avg_window}"
+        )
+        screen.blit(font.render(header, True, TEXT), (10, 8))
+        screen.blit(
+            small_font.render("mouse wheel: zoom horizontal axis", True, AXIS),
+            (10, height - 24),
+        )
+
+        legend_x = 10
+        for trace in TRACE_DEFS:
+            pygame.draw.line(
+                screen,
+                trace["color"],
+                (legend_x, 43),
+                (legend_x + 24, 43),
+                3,
+            )
+            screen.blit(small_font.render(trace["label"], True, TEXT), (legend_x + 30, 34))
+            legend_x += 120
+
+        if error_message:
+            error_surface = font.render(error_message, True, ERROR)
+            screen.blit(error_surface, (10, 56))
+
+        pygame.display.flip()
+        clock.tick(60)
+
+    pygame.quit()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())