icassp2027-experiment/test_easycom_oracle_features.py

"""Focused tests for the causal EasyCom oracle-information experiment."""

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace
import tempfile
import unittest

import numpy as np
import torch

from easycom_head import TurnPredictionHead
from local_turn_features import ORACLE_FORMAT, load_frame_set_features, load_local_features
from prepare_easycom_oracle_features import (
    causal_history,
    oracle_person_rows,
    participant_activity,
    participant_direct_addressee_activity,
)


class OracleFeatureTests(unittest.TestCase):
    def setUp(self) -> None:
        torch.set_num_threads(1)

    def test_transcript_intervals_match_egospeak_right_edge_sampling(self) -> None:
        rows = [
            SimpleNamespace(
                uid="a",
                participant_id=3,
                # EasyCom annotation Start_Frame=1, End_Frame=8. Samples at
                # source markers 4,8,... therefore include only row 0.
                dense_start_frame=0,
                dense_end_frame_exclusive=8,
            ),
            SimpleNamespace(
                uid="b",
                participant_id=2,
                # Start_Frame=5, End_Frame=13 includes markers 8 and 12.
                dense_start_frame=4,
                dense_end_frame_exclusive=13,
            ),
        ]
        activity, wearer, audit = participant_activity(
            rows, num_steps=4, wearer_id=2
        )
        self.assertEqual(activity[3].tolist(), [True, False, False, False])
        self.assertEqual(wearer.tolist(), [False, True, True, False])
        self.assertEqual(
            audit["target_sampling_policy"],
            "egospeak-right-edge-start-inclusive-end-exclusive",
        )

    def test_history_at_row_does_not_read_later_activity(self) -> None:
        first = torch.tensor([0, 1, 0, 0, 0], dtype=torch.bool)
        second = first.clone()
        second[3:] = True
        observed = causal_history(first, row=2, history_steps=4)
        changed_future = causal_history(second, row=2, history_steps=4)
        self.assertTrue(torch.equal(observed, changed_future))
        self.assertEqual(
            observed.tolist(),
            [0, 0, 1, 0, 0, 1, 1, 1],
        )
        delayed = causal_history(second, row=2, history_steps=4, delay_steps=2)
        self.assertEqual(
            delayed.tolist(),
            [0, 0, 0, 0, 0, 0, 0, 1],
        )
        visibility = torch.tensor([1, 0, 1, 1, 1], dtype=torch.bool)
        visible_only = causal_history(
            second,
            row=2,
            history_steps=4,
            observed_sequence=visibility,
        )
        self.assertEqual(
            visible_only.tolist(),
            [0, 0, 0, 0, 0, 1, 0, 1],
        )

    def test_direct_addressee_keeps_group_and_unknown_out_of_positive_channel(self) -> None:
        rows = [
            SimpleNamespace(
                uid="direct",
                participant_id=3,
                target_of_speech=(2,),
                dense_start_frame=0,
                dense_end_frame_exclusive=8,
            ),
            SimpleNamespace(
                uid="group",
                participant_id=4,
                target_of_speech=("Group",),
                dense_start_frame=4,
                dense_end_frame_exclusive=13,
            ),
            SimpleNamespace(
                uid="unknown",
                participant_id=5,
                target_of_speech=("None",),
                dense_start_frame=8,
                dense_end_frame_exclusive=17,
            ),
        ]
        directed, audit = participant_direct_addressee_activity(
            rows, num_steps=5, wearer_id=2
        )
        self.assertEqual(directed[3].tolist(), [True, False, False, False, False])
        self.assertNotIn(4, directed)
        self.assertNotIn(5, directed)
        self.assertEqual(
            audit["target_of_speech_utterance_counts"],
            {"direct_wearer": 1, "group": 1, "other": 0, "unknown": 1},
        )
        self.assertTrue(audit["target_of_speech_group_unknown_not_coerced"])

    def test_addressee_oracle_has_capacity_matched_zero_control(self) -> None:
        activity = {3: torch.tensor([1, 1, 0], dtype=torch.bool)}
        addressed = {3: torch.tensor([1, 0, 0], dtype=torch.bool)}
        visibility = {3: torch.ones(3, dtype=torch.bool)}
        visible = [[3], [3], [3]]
        records = [
            [{"id": "3", "source_clip_index": 0, "source_time": row / 5}]
            for row in range(3)
        ]
        control, control_offsets, _, _ = oracle_person_rows(
            activity,
            visibility,
            visible,
            records,
            condition="aggregate-visible-address-control",
            history_steps=2,
            delay_steps=0,
            seed=0,
            direct_addressee_activity=addressed,
        )
        oracle, oracle_offsets, _, _ = oracle_person_rows(
            activity,
            visibility,
            visible,
            records,
            condition="aggregate-visible-addressee",
            history_steps=2,
            delay_steps=0,
            seed=0,
            direct_addressee_activity=addressed,
        )
        self.assertEqual(control.shape, (3, 6))
        self.assertTrue(torch.equal(control_offsets, oracle_offsets))
        self.assertTrue(torch.equal(control[:, :4], oracle[:, :4]))
        self.assertEqual(control[:, 4:].tolist(), [[0, 0], [0, 0], [0, 0]])
        self.assertEqual(oracle[:, 4:].tolist(), [[0, 1], [1, 0], [0, 0]])

    def test_shuffled_id_preserves_every_lag_count_but_breaks_histories(self) -> None:
        activity = {
            3: torch.tensor([1, 1, 0, 0, 1, 1], dtype=torch.bool),
            4: torch.tensor([0, 0, 1, 1, 0, 0], dtype=torch.bool),
            5: torch.tensor([0, 1, 0, 1, 0, 1], dtype=torch.bool),
        }
        visible = [[3, 4, 5] for _ in range(6)]
        visibility = {
            person_id: torch.ones(6, dtype=torch.bool)
            for person_id in activity
        }
        records = [
            [
                {"id": str(person_id), "source_clip_index": 0, "source_time": row / 5}
                for person_id in person_ids
            ]
            for row, person_ids in enumerate(visible)
        ]
        original, offsets, _, _ = oracle_person_rows(
            activity,
            visibility,
            visible,
            records,
            condition="per-person",
            history_steps=4,
            delay_steps=5,
            seed=9,
        )
        shuffled, shuffled_offsets, _, audit = oracle_person_rows(
            activity,
            visibility,
            visible,
            records,
            condition="shuffled-id",
            history_steps=4,
            delay_steps=5,
            seed=9,
        )
        self.assertTrue(torch.equal(offsets, shuffled_offsets))
        for row in range(6):
            start, stop = int(offsets[row]), int(offsets[row + 1])
            self.assertTrue(
                torch.equal(
                    original[start:stop, :4].sum(dim=0),
                    shuffled[start:stop, :4].sum(dim=0),
                )
            )
        self.assertGreater(audit["shuffled_values_changed"], 0)

    def test_oracle_format_requires_explicit_loader_and_head_mode(self) -> None:
        n = 4
        times = (torch.arange(n, dtype=torch.float64) + 1) / 5
        reference = {
            "states": torch.randn(n, 8),
            "label_indices": torch.arange(n),
            "times": times,
            "mode": "open-stream-replay",
            "start": 0,
        }
        sidecar = {
            "format": ORACLE_FORMAT,
            "dataset": "easycom",
            "sample_id": "Session_4",
            "times": times.clone(),
            "label_indices": torch.arange(n),
            "partial": False,
            "visual": torch.randn(n, 6),
            "audio": torch.randn(n, 5),
            "visual_valid": torch.ones(n, dtype=torch.bool),
            "audio_valid": torch.ones(n, dtype=torch.bool),
            "persons": {
                "features": torch.randn(n, 6),
                # Oracle person rows are current-state observations sampled at
                # the same inclusive 5 Hz right edges as the target timeline.
                "frame_times": times.clone(),
                "frame_offsets": torch.arange(n + 1, dtype=torch.long),
                "records": [{"id": "3"} for _ in range(n)],
            },
            "config": {
                "oracle_analysis": True,
                "uses_speech_labels": True,
                "causal_policy": "egospeak-right-edge-current-state",
                "oracle_condition": "per-person",
                "visual_dim": 6,
                "audio_dim": 5,
            },
        }
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            torch.save(sidecar, path)
            with self.assertRaisesRegex(ValueError, "format"):
                load_local_features(
                    path,
                    reference,
                    sample_id="Session_4",
                    dataset="easycom",
                    modalities="av",
                )
            audio, people, contract = load_frame_set_features(
                path,
                reference,
                sample_id="Session_4",
                dataset="easycom",
                oracle_analysis=True,
            )
        self.assertEqual(audio.shape, (n, 6))
        self.assertEqual(people.dimension, 6)
        self.assertEqual(contract["format"], ORACLE_FORMAT)
        self.assertTrue(torch.equal(people.row_to_frame, torch.arange(n)))
        self.assertTrue(torch.allclose(people.chosen_frame_times, times))
        head = TurnPredictionHead(
            hidden_size=8,
            num_horizons=2,
            num_classes=3,
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=4,
            local_fusion="oracle-frame-set-attention",
            person_attention_dim=4,
            person_attention_heads=1,
        )
        padded, mask = people.padded()
        logits = head(reference["states"], audio, padded, mask)
        self.assertEqual(logits.shape, (n, 2, 3))


if __name__ == "__main__":
    unittest.main()