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

tests/test_models.py

@@ -0,0 +1,69 @@
+"""Tests for TrajectoryFlowModel architecture."""
+from __future__ import annotations
+
+import torch
+import pytest
+
+from ai_mouse.models import TrajectoryFlowModel
+
+
+class TestTrajectoryFlowModel:
+    """Test the Conditional Flow Matching model."""
+
+    @pytest.fixture
+    def model(self):
+        return TrajectoryFlowModel(
+            seq_len=64, d_model=128, nhead=4, num_layers=4,
+            dim_feedforward=256, dropout=0.1, cond_dim=3,
+        )
+
+    def test_output_shape(self, model):
+        """(4, 64, 3) input → (4, 64, 3) output."""
+        batch = 4
+        x_t = torch.randn(batch, 64, 3)
+        t = torch.rand(batch)
+        cond = torch.randn(batch, 3)
+        out = model(x_t, t, cond)
+        assert out.shape == (batch, 64, 3)
+
+    def test_single_sample(self, model):
+        """(1, 64, 3) works."""
+        x_t = torch.randn(1, 64, 3)
+        t = torch.rand(1)
+        cond = torch.randn(1, 3)
+        out = model(x_t, t, cond)
+        assert out.shape == (1, 64, 3)
+
+    def test_deterministic(self, model):
+        """Eval mode, same input → same output."""
+        model.eval()
+        x_t = torch.randn(2, 64, 3)
+        t = torch.tensor([0.3, 0.7])
+        cond = torch.randn(2, 3)
+        with torch.no_grad():
+            out1 = model(x_t, t, cond)
+            out2 = model(x_t, t, cond)
+        torch.testing.assert_close(out1, out2)
+
+    def test_different_timesteps(self, model):
+        """t=0.1 vs t=0.9 gives different output."""
+        model.eval()
+        x_t = torch.randn(1, 64, 3)
+        cond = torch.randn(1, 3)
+        with torch.no_grad():
+            out_early = model(x_t, torch.tensor([0.1]), cond)
+            out_late = model(x_t, torch.tensor([0.9]), cond)
+        assert not torch.allclose(out_early, out_late, atol=1e-5)
+
+    def test_gradient_flows(self, model):
+        """Backward works, grad on x_t exists."""
+        model.train()
+        x_t = torch.randn(2, 64, 3, requires_grad=True)
+        t = torch.rand(2)
+        cond = torch.randn(2, 3)
+        out = model(x_t, t, cond)
+        loss = out.sum()
+        loss.backward()
+        assert x_t.grad is not None
+        assert x_t.grad.shape == (2, 64, 3)
+        assert x_t.grad.abs().sum() > 0