From 17b406c28b74fc9ea46e3f52d67ae3c5e3a51e15 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Tue, 12 May 2026 01:02:25 +0800
Subject: [PATCH] feat(lib): add private _coord.py with numpy transforms

Copy of coord.py (which is already pure numpy) into the private
underscored module to be consumed by upcoming mouse.py rewrite.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 src/ai_mouse/_coord.py    | 81 +++++++++++++++++++++++++++++++++++++++
 tests/unit/test__coord.py | 30 +++++++++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 src/ai_mouse/_coord.py
 create mode 100644 tests/unit/test__coord.py

diff --git a/src/ai_mouse/_coord.py b/src/ai_mouse/_coord.py
new file mode 100644
index 0000000..ff3d528
--- /dev/null
+++ b/src/ai_mouse/_coord.py
@@ -0,0 +1,81 @@
+"""Rotated coordinate system for angle-invariant trajectory encoding.
+
+All trajectories are normalised into a frame where:
+- start → (0, 0)
+- end   → (1, 0)
+- lateral displacement is perpendicular to start→end axis
+
+This makes the model angle-invariant: a 45° diagonal move and a horizontal
+move look identical in the rotated frame (just "forward from 0 to 1").
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+
+def encode_trajectory(
+    points: np.ndarray,
+    start: tuple[int, int],
+    end: tuple[int, int],
+) -> np.ndarray:
+    """Transform pixel coordinates to rotated normalised frame.
+
+    Args:
+        points: (N, 2) array of (x, y) pixel coordinates.
+        start: (x, y) start position.
+        end: (x, y) end position.
+
+    Returns:
+        (N, 2) array of (forward, lateral) in normalised rotated frame.
+    """
+    sx, sy = float(start[0]), float(start[1])
+    ex, ey = float(end[0]), float(end[1])
+    dist = math.hypot(ex - sx, ey - sy)
+    if dist < 1e-8:
+        return np.zeros_like(points)
+
+    ux, uy = (ex - sx) / dist, (ey - sy) / dist
+    vx, vy = -uy, ux
+
+    dx = points[:, 0] - sx
+    dy = points[:, 1] - sy
+
+    forward = (dx * ux + dy * uy) / dist
+    lateral = (dx * vx + dy * vy) / dist
+
+    return np.stack([forward, lateral], axis=1)
+
+
+def decode_trajectory(
+    normalised: np.ndarray,
+    start: tuple[int, int],
+    end: tuple[int, int],
+) -> np.ndarray:
+    """Transform rotated normalised frame back to pixel coordinates.
+
+    Args:
+        normalised: (N, 2) array of (forward, lateral).
+        start: (x, y) start position.
+        end: (x, y) end position.
+
+    Returns:
+        (N, 2) array of (x, y) pixel coordinates.
+    """
+    sx, sy = float(start[0]), float(start[1])
+    ex, ey = float(end[0]), float(end[1])
+    dist = math.hypot(ex - sx, ey - sy)
+    if dist < 1e-8:
+        return np.full_like(normalised, [sx, sy])
+
+    ux, uy = (ex - sx) / dist, (ey - sy) / dist
+    vx, vy = -uy, ux
+
+    forward = normalised[:, 0]
+    lateral = normalised[:, 1]
+
+    px = sx + forward * dist * ux + lateral * dist * vx
+    py = sy + forward * dist * uy + lateral * dist * vy
+
+    return np.stack([px, py], axis=1)
diff --git a/tests/unit/test__coord.py b/tests/unit/test__coord.py
new file mode 100644
index 0000000..ab8b422
--- /dev/null
+++ b/tests/unit/test__coord.py
@@ -0,0 +1,30 @@
+"""Test the private numpy coordinate transforms."""
+from __future__ import annotations
+
+import numpy as np
+
+from ai_mouse._coord import decode_trajectory, encode_trajectory
+
+
+def test_encode_decode_roundtrip() -> None:
+    points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]])
+    start = (100, 200)
+    end = (500, 300)
+    encoded = encode_trajectory(points, start, end)
+    decoded = decode_trajectory(encoded, start, end)
+    assert np.allclose(decoded, points, atol=1e-6)
+
+
+def test_encode_endpoints() -> None:
+    """Start should encode to (0,0); end should encode to (1,0)."""
+    points = np.array([[100.0, 200.0], [500.0, 300.0]])
+    encoded = encode_trajectory(points, (100, 200), (500, 300))
+    assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6)
+    assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6)
+
+
+def test_zero_distance_returns_zeros() -> None:
+    points = np.array([[100.0, 200.0]])
+    encoded = encode_trajectory(points, (100, 200), (100, 200))
+    assert encoded.shape == (1, 2)
+    assert np.all(encoded == 0)