51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
|
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()
|