173 lines
6.2 KiB
Python
173 lines
6.2 KiB
Python
"""Tests for Flow Matching training pipeline."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from ai_mouse.trainer import load_and_prepare_data, train, _augment
|
|
|
|
|
|
def _make_synthetic_trace(start, end, n_moves=30):
|
|
"""Create a synthetic trace dict mimicking real JSONL format."""
|
|
sx, sy = start
|
|
ex, ey = end
|
|
dist = math.hypot(ex - sx, ey - sy)
|
|
angle = math.degrees(math.atan2(ey - sy, ex - sx))
|
|
|
|
events = []
|
|
for i in range(n_moves):
|
|
t_frac = i / (n_moves - 1)
|
|
x = int(sx + (ex - sx) * t_frac + np.random.normal(0, 3))
|
|
y = int(sy + (ey - sy) * t_frac + np.random.normal(0, 3))
|
|
t = int(t_frac * 500 + np.random.normal(0, 5))
|
|
events.append({"type": "move", "x": x, "y": y, "t": max(0, t)})
|
|
|
|
events.sort(key=lambda e: e["t"])
|
|
events[0]["t"] = 0
|
|
|
|
last_t = events[-1]["t"]
|
|
events.append({"type": "down", "x": ex, "y": ey, "t": last_t + 50})
|
|
events.append({"type": "up", "x": ex, "y": ey, "t": last_t + 130})
|
|
|
|
return {
|
|
"meta": {"start": [sx, sy], "end": [ex, ey], "dist": int(dist), "angle": round(angle, 1)},
|
|
"events": events,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def synthetic_traces_file(tmp_path):
|
|
"""Create a temp JSONL file with 25 synthetic traces."""
|
|
traces_path = tmp_path / "traces.jsonl"
|
|
rng = np.random.default_rng(42)
|
|
lines = []
|
|
for _ in range(25):
|
|
sx, sy = int(rng.integers(50, 750)), int(rng.integers(50, 750))
|
|
angle = rng.uniform(0, 2 * math.pi)
|
|
dist = int(rng.integers(100, 500))
|
|
ex = int(sx + dist * math.cos(angle))
|
|
ey = int(sy + dist * math.sin(angle))
|
|
ex = max(0, min(800, ex))
|
|
ey = max(0, min(600, ey))
|
|
trace = _make_synthetic_trace((sx, sy), (ex, ey))
|
|
lines.append(json.dumps(trace))
|
|
traces_path.write_text("\n".join(lines), encoding="utf-8")
|
|
return traces_path
|
|
|
|
|
|
class TestLoadAndPrepare:
|
|
def test_returns_correct_shapes(self, synthetic_traces_file):
|
|
seq, cond, click_durs = load_and_prepare_data(synthetic_traces_file, seq_len=64)
|
|
assert seq.shape[1] == 64
|
|
assert seq.shape[2] == 3
|
|
assert cond.shape[1] == 3
|
|
assert len(seq) > 0
|
|
|
|
def test_forward_starts_near_zero(self, synthetic_traces_file):
|
|
seq, cond, _ = load_and_prepare_data(synthetic_traces_file, seq_len=64)
|
|
assert abs(seq[:, 0, 0].mean()) < 0.15
|
|
|
|
def test_forward_ends_near_one(self, synthetic_traces_file):
|
|
seq, cond, _ = load_and_prepare_data(synthetic_traces_file, seq_len=64)
|
|
assert abs(seq[:, -1, 0].mean() - 1.0) < 0.15
|
|
|
|
|
|
class TestAugment:
|
|
def test_augmentation_multiplies_data(self, synthetic_traces_file):
|
|
seq, cond, _ = load_and_prepare_data(synthetic_traces_file, seq_len=64)
|
|
n_orig = len(seq)
|
|
seq_aug, cond_aug = _augment(seq, cond)
|
|
assert len(seq_aug) == n_orig * 6
|
|
assert len(cond_aug) == n_orig * 6
|
|
|
|
|
|
class TestTrain:
|
|
def test_train_produces_model_files(self, synthetic_traces_file, tmp_path):
|
|
output_dir = tmp_path / "models"
|
|
train(
|
|
data_path=synthetic_traces_file,
|
|
output_dir=output_dir,
|
|
epochs=3,
|
|
batch_size=8,
|
|
seq_len=64,
|
|
)
|
|
assert (output_dir / "flow_model.pt").exists()
|
|
assert (output_dir / "click_dist.json").exists()
|
|
assert (output_dir / "duration_dist.json").exists()
|
|
assert (output_dir / "train_config.json").exists()
|
|
|
|
def test_train_loss_decreases(self, synthetic_traces_file, tmp_path):
|
|
output_dir = tmp_path / "models"
|
|
losses = []
|
|
|
|
def cb(msg):
|
|
if "loss" in msg:
|
|
losses.append(msg["loss"])
|
|
|
|
train(
|
|
data_path=synthetic_traces_file,
|
|
output_dir=output_dir,
|
|
epochs=20,
|
|
batch_size=8,
|
|
seq_len=64,
|
|
progress_callback=cb,
|
|
)
|
|
first_half = np.mean(losses[:10])
|
|
second_half = np.mean(losses[10:])
|
|
assert second_half < first_half
|
|
|
|
|
|
class TestTrajectoryDataset:
|
|
def test_dataset_length_with_augmentation(self):
|
|
"""Dataset length = N * 6 when augment=True."""
|
|
from ai_mouse.trainer import TrajectoryDataset
|
|
seq = np.zeros((10, 64, 3), dtype=np.float32)
|
|
cond = np.zeros((10, 3), dtype=np.float32)
|
|
ds = TrajectoryDataset(seq, cond, augment=True)
|
|
assert len(ds) == 60
|
|
|
|
def test_dataset_length_without_augmentation(self):
|
|
from ai_mouse.trainer import TrajectoryDataset
|
|
seq = np.zeros((10, 64, 3), dtype=np.float32)
|
|
cond = np.zeros((10, 3), dtype=np.float32)
|
|
ds = TrajectoryDataset(seq, cond, augment=False)
|
|
assert len(ds) == 10
|
|
|
|
def test_getitem_returns_tensors(self):
|
|
from ai_mouse.trainer import TrajectoryDataset
|
|
import torch
|
|
seq = np.random.randn(5, 64, 3).astype(np.float32)
|
|
cond = np.random.randn(5, 3).astype(np.float32)
|
|
ds = TrajectoryDataset(seq, cond, augment=True)
|
|
s, c = ds[0]
|
|
assert isinstance(s, torch.Tensor)
|
|
assert isinstance(c, torch.Tensor)
|
|
assert s.shape == (64, 3)
|
|
assert c.shape == (3,)
|
|
|
|
def test_aug_id_zero_returns_original(self):
|
|
"""Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged."""
|
|
from ai_mouse.trainer import TrajectoryDataset
|
|
import torch
|
|
seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32)
|
|
cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32)
|
|
ds = TrajectoryDataset(seq, cond, augment=True)
|
|
s0, c0 = ds[0]
|
|
np.testing.assert_allclose(s0.numpy(), seq[0], rtol=1e-5)
|
|
np.testing.assert_allclose(c0.numpy(), cond[0], rtol=1e-5)
|
|
|
|
def test_aug_id_one_flips_lateral(self):
|
|
"""Aug id 1 should flip the sign of the lateral channel (index 1)."""
|
|
from ai_mouse.trainer import TrajectoryDataset
|
|
seq = np.zeros((1, 64, 3), dtype=np.float32)
|
|
seq[0, :, 1] = 0.5 # lateral all positive
|
|
cond = np.zeros((1, 3), dtype=np.float32)
|
|
ds = TrajectoryDataset(seq, cond, augment=True)
|
|
# idx=1 → base_idx=0, aug_id=1 → flip
|
|
s1, _ = ds[1]
|
|
assert (s1[:, 1] < 0).all()
|