Add floor-event timing objective and addressee diagnostics

test_easycom_floor_event_probe.py

@@ -0,0 +1,120 @@
+"""Focused tests for the direct competing-risk floor-event probe."""
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import torch
+
+from train_easycom_floor_event_probe import (
+    DirectFloorEventHead,
+    FloorEventDataset,
+    ProbeSession,
+    metric_rows,
+)
+
+
+def activity(wearer: int, others: int) -> list[int]:
+    return [int(not wearer and not others), wearer, others]
+
+
+def probe_session() -> ProbeSession:
+    target = np.asarray(
+        [
+            activity(0, 1),
+            activity(0, 0),
+            activity(1, 0),
+            activity(1, 0),
+            activity(1, 1),
+            activity(0, 1),
+            activity(0, 1),
+            activity(0, 0),
+            activity(0, 0),
+            activity(0, 1),
+        ],
+        dtype=np.float32,
+    )
+    steps = len(target)
+    return ProbeSession(
+        name="Session_12",
+        states=torch.randn(steps, 2, 4),
+        audio=torch.cat((torch.randn(steps, 3), torch.ones(steps, 1)), dim=1),
+        oracle=torch.randn(steps, 6),
+        oracle_valid=torch.ones(steps, dtype=torch.bool),
+        target=target,
+        oracle_condition="aggregate-visible",
+        oracle_contract={"config": {"oracle_condition": "aggregate-visible"}},
+    )
+
+
+class FloorEventProbeTests(unittest.TestCase):
+    def setUp(self) -> None:
+        torch.set_num_threads(1)
+
+    def test_dataset_uses_requested_decision_source(self) -> None:
+        session = probe_session()
+        pre = FloorEventDataset(
+            [session], decision_timing="pre-boundary", midturn_pause_max_steps=10
+        )
+        post = FloorEventDataset(
+            [session], decision_timing="post-boundary", midturn_pause_max_steps=10
+        )
+        pre_response = next(
+            item
+            for item in pre.examples
+            if item.event_index == 0 and item.horizon_index == 1 and item.label
+        )
+        post_response = next(
+            item
+            for item in post.examples
+            if item.event_index == 0 and item.horizon_index == 1 and item.label
+        )
+        self.assertEqual(pre_response.source_row, 0)
+        self.assertEqual(post_response.source_row, 1)
+
+    def test_head_outputs_two_events_by_ten_horizons(self) -> None:
+        model = DirectFloorEventHead(
+            state_shape=(2, 4),
+            audio_dim=3,
+            oracle_dim=6,
+            layer_fusion="concat-mlp",
+            local_projection_dim=4,
+            person_attention_dim=4,
+            person_attention_heads=1,
+            dropout=0.0,
+        )
+        output = model(
+            torch.randn(5, 2, 4),
+            torch.cat((torch.randn(5, 3), torch.ones(5, 1)), dim=1),
+            torch.randn(5, 6),
+            torch.ones(5),
+        )
+        self.assertEqual(output.shape, (5, 2, 10))
+
+    def test_fixed_threshold_is_applied_without_reselection(self) -> None:
+        predictions = {
+            "score": np.tile(np.asarray([0.9, 0.4]), 20),
+            "event": np.repeat(np.arange(2), 20),
+            "horizon": np.tile(np.repeat(np.arange(10), 2), 2),
+            "label": np.tile(np.asarray([True, False]), 20),
+            "short_resume": np.zeros(40, dtype=bool),
+        }
+        thresholds = {
+            (event, horizon): 0.95
+            for event in range(2)
+            for horizon in range(10)
+        }
+        rows, _ = metric_rows(
+            predictions,
+            matched_recall=0.8,
+            fixed_thresholds=thresholds,
+            threshold_source="Session_12",
+        )
+        self.assertEqual(len(rows), 20)
+        self.assertEqual(rows[0]["threshold"], 0.95)
+        self.assertEqual(rows[0]["achieved_positive_recall"], 0.0)
+
+
+if __name__ == "__main__":
+    unittest.main()