chore: initialize git repo, add matplotlib dep, extend config
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
This commit is contained in:
106
tests/test_generator.py
Normal file
106
tests/test_generator.py
Normal file
@@ -0,0 +1,106 @@
|
||||
"""Tests for Flow Matching trajectory generator."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from ai_mouse.generator import generate
|
||||
from ai_mouse.models import TrajectoryFlowModel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_dir(tmp_path):
|
||||
"""Create temp dir with Flow model artifacts."""
|
||||
model = TrajectoryFlowModel(seq_len=64)
|
||||
torch.save(model.state_dict(), tmp_path / "flow_model.pt")
|
||||
|
||||
click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
|
||||
(tmp_path / "click_dist.json").write_text(json.dumps(click_dist))
|
||||
|
||||
duration_dist = {
|
||||
"bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
|
||||
"params": [
|
||||
{"mu_log": 5.5, "sigma_log": 0.3},
|
||||
{"mu_log": 5.8, "sigma_log": 0.3},
|
||||
{"mu_log": 6.0, "sigma_log": 0.3},
|
||||
{"mu_log": 6.2, "sigma_log": 0.3},
|
||||
{"mu_log": 6.5, "sigma_log": 0.3},
|
||||
{"mu_log": 6.7, "sigma_log": 0.3},
|
||||
{"mu_log": 6.9, "sigma_log": 0.3},
|
||||
{"mu_log": 7.0, "sigma_log": 0.3},
|
||||
],
|
||||
}
|
||||
(tmp_path / "duration_dist.json").write_text(json.dumps(duration_dist))
|
||||
|
||||
train_config = {
|
||||
"seq_len": 64,
|
||||
"d_model": 128,
|
||||
"nhead": 4,
|
||||
"num_layers": 4,
|
||||
"dim_feedforward": 256,
|
||||
"cond_dim": 3,
|
||||
}
|
||||
(tmp_path / "train_config.json").write_text(json.dumps(train_config))
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestGenerate:
|
||||
def test_returns_list_of_tuples(self, model_dir):
|
||||
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
assert isinstance(result, list)
|
||||
assert all(isinstance(p, tuple) and len(p) == 3 for p in result)
|
||||
# All elements are ints
|
||||
for p in result:
|
||||
assert all(isinstance(v, int) for v in p)
|
||||
|
||||
def test_timestamps_monotonically_increasing(self, model_dir):
|
||||
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
times = [p[2] for p in result]
|
||||
for i in range(1, len(times)):
|
||||
assert times[i] >= times[i - 1]
|
||||
|
||||
def test_starts_near_start(self, model_dir):
|
||||
start = (100, 200)
|
||||
result = generate(start=start, end=(500, 400), model_dir=str(model_dir))
|
||||
first = result[0]
|
||||
assert abs(first[0] - start[0]) < 30
|
||||
assert abs(first[1] - start[1]) < 30
|
||||
|
||||
def test_ends_near_end(self, model_dir):
|
||||
end = (500, 400)
|
||||
result = generate(start=(100, 200), end=end, model_dir=str(model_dir))
|
||||
# Last two are click events; the one before is last movement point
|
||||
last_move = result[-3]
|
||||
assert abs(last_move[0] - end[0]) < 30
|
||||
assert abs(last_move[1] - end[1]) < 30
|
||||
|
||||
def test_last_two_are_click_events(self, model_dir):
|
||||
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
down = result[-2]
|
||||
up = result[-1]
|
||||
# Same x, y for click down and up
|
||||
assert down[0] == up[0]
|
||||
assert down[1] == up[1]
|
||||
# Up timestamp > down timestamp
|
||||
assert up[2] > down[2]
|
||||
# Click duration within bounds
|
||||
assert 20 <= up[2] - down[2] <= 300
|
||||
|
||||
def test_different_z_gives_different_paths(self, model_dir):
|
||||
r1 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
r2 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
points1 = [(p[0], p[1]) for p in r1[:-2]]
|
||||
points2 = [(p[0], p[1]) for p in r2[:-2]]
|
||||
assert points1 != points2
|
||||
|
||||
def test_n_points_parameter(self, model_dir):
|
||||
result = generate(
|
||||
start=(100, 200), end=(500, 400), n_points=32, model_dir=str(model_dir)
|
||||
)
|
||||
# 32 move points + 2 click events = 34
|
||||
assert len(result) == 34
|
||||
Reference in New Issue
Block a user