Add EgoSpeak prediction alignment for floor evaluation

test_prepare_egospeak_floor_eval.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+
+from prepare_egospeak_floor_eval import (
+    NUM_CLASSES,
+    NUM_HORIZONS,
+    align_miniroad_prediction,
+)
+
+
+class EgoSpeakFloorAlignmentTests(unittest.TestCase):
+    def test_source_forecasts_are_aligned_to_target_rows(self) -> None:
+        target_steps = 14
+        raw = np.zeros(
+            (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES),
+            dtype=np.float32,
+        )
+        for source_row in range(len(raw)):
+            for horizon_index in range(NUM_HORIZONS):
+                raw[source_row, horizon_index, 1] = (
+                    100 * source_row + horizon_index
+                ) / 1000
+
+        aligned, valid = align_miniroad_prediction(raw, target_steps)
+        self.assertEqual(aligned.shape, (target_steps, NUM_CLASSES, NUM_HORIZONS))
+        self.assertEqual(valid.shape, (target_steps, NUM_HORIZONS))
+        for source_row in range(len(raw)):
+            for horizon_index in range(NUM_HORIZONS):
+                target_row = source_row + horizon_index + 1
+                self.assertEqual(
+                    aligned[target_row, 1, horizon_index],
+                    raw[source_row, horizon_index, 1],
+                )
+                self.assertTrue(valid[target_row, horizon_index])
+        for horizon_index in range(NUM_HORIZONS):
+            self.assertFalse(valid[: horizon_index + 1, horizon_index].any())
+
+    def test_shape_mismatch_is_rejected(self) -> None:
+        with self.assertRaisesRegex(ValueError, "Expected MiniROADA prediction"):
+            align_miniroad_prediction(
+                np.zeros((4, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32),
+                14,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()