"""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() class TestResumeFrom: def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path): """train() with resume_from should load weights from given checkpoint dir.""" import torch from ai_mouse.trainer import train from ai_mouse.models import TrajectoryFlowModel # First, train an initial model and save it ckpt_dir = tmp_path / "pretrain" train( data_path=synthetic_traces_file, output_dir=ckpt_dir, epochs=2, batch_size=8, seq_len=64, ) assert (ckpt_dir / "flow_model.pt").exists() # Read its weights to compare later m_pretrain = TrajectoryFlowModel(seq_len=64) m_pretrain.load_state_dict(torch.load(ckpt_dir / "flow_model.pt", weights_only=True)) first_param_pre = next(m_pretrain.parameters()).clone() # Now train with resume_from for 1 epoch — weights should still be loaded out_dir = tmp_path / "finetune" train( data_path=synthetic_traces_file, output_dir=out_dir, epochs=1, batch_size=8, seq_len=64, resume_from=ckpt_dir, ) m_after = TrajectoryFlowModel(seq_len=64) m_after.load_state_dict(torch.load(out_dir / "flow_model.pt", weights_only=True)) first_param_after = next(m_after.parameters()) # After 1 epoch, weights should be close to pre-train, not random init # (random init would be O(1) magnitude apart; 1 epoch on small data shifts O(0.1)) diff = (first_param_pre - first_param_after).abs().mean().item() assert diff < 0.5, f"Resume_from weights diverged too much: {diff}" def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path): from ai_mouse.trainer import train with pytest.raises(FileNotFoundError): train( data_path=synthetic_traces_file, output_dir=tmp_path / "out", epochs=1, batch_size=8, seq_len=64, resume_from=tmp_path / "nonexistent", )