tests/test_chart_features.py

from __future__ import annotations

import numpy as np
import pytest

from cil.chart_features import (
    CONTEXT_HASH_WIDTH,
    OBSERVATION_EMBED_DIM,
    OBJECT_LAYOUT_EMBED_DIM,
    build_chart_feature,
    chart_feature_dim,
)


def test_base_chart_feature_preserves_original_flattened_base_action() -> None:
    base_action = np.arange(12, dtype=np.float32).reshape(3, 4)

    feature = build_chart_feature(base_action, mode="base")

    assert feature.dtype == np.float32
    assert np.allclose(feature, base_action.reshape(-1))
    assert chart_feature_dim(base_action, mode="base") == 12


def test_base_context_feature_is_deterministic_and_extends_base() -> None:
    base_action = np.zeros((2, 7), dtype=np.float32)
    metadata = {"task_id": "PickCube-v1", "instruction": "pick up the cube"}

    first = build_chart_feature(base_action, metadata, mode="base_context")
    second = build_chart_feature(base_action, metadata, mode="base_context")

    assert np.allclose(first, second)
    assert first.shape == (base_action.size + 2 * CONTEXT_HASH_WIDTH + 2,)
    assert chart_feature_dim(base_action, mode="base_context") == first.shape[0]


def test_base_context_feature_does_not_read_outcome_or_hidden_fields() -> None:
    base_action = np.ones((1, 4), dtype=np.float32)
    shared = {"task_id": "StackCube-v1", "instruction": "stack the cube"}
    metadata_a = {
        **shared,
        "scalar_utility": 1.0,
        "delta_U": 0.8,
        "label": "positive",
        "outcome_vector": {"success": True},
        "hidden_chart_oracle": 1.0,
    }
    metadata_b = {
        **shared,
        "scalar_utility": -1.0,
        "delta_U": -0.4,
        "label": "negative",
        "outcome_vector": {"success": False},
        "hidden_chart_oracle": 0.0,
    }

    assert np.allclose(
        build_chart_feature(base_action, metadata_a, mode="base_context"),
        build_chart_feature(base_action, metadata_b, mode="base_context"),
    )


def test_base_context_obs_appends_precomputed_observation_embedding(tmp_path) -> None:
    base_action = np.zeros((2, 7), dtype=np.float32)
    embeddings = np.arange(OBSERVATION_EMBED_DIM * 2, dtype=np.float32).reshape(
        2, OBSERVATION_EMBED_DIM
    )
    path = tmp_path / "obs_embeddings.npz"
    np.savez_compressed(path, embeddings=embeddings)
    metadata = {
        "task_id": "PullCube-v1",
        "instruction": "pull the cube",
        "observation_embedding_path": f"{path.name}#embeddings/1",
        "_chart_root": str(tmp_path),
        "scalar_utility": -999.0,
        "label": "negative",
    }

    feature = build_chart_feature(base_action, metadata, mode="base_context_obs")

    assert feature.shape == (
        base_action.size + 2 * CONTEXT_HASH_WIDTH + 2 + OBSERVATION_EMBED_DIM,
    )
    assert np.allclose(feature[-OBSERVATION_EMBED_DIM:], embeddings[1])
    assert chart_feature_dim(base_action, mode="base_context_obs") == feature.shape[0]


def test_base_context_obj_appends_precomputed_object_embedding(tmp_path) -> None:
    base_action = np.zeros((2, 7), dtype=np.float32)
    embeddings = np.arange(OBJECT_LAYOUT_EMBED_DIM * 2, dtype=np.float32).reshape(
        2, OBJECT_LAYOUT_EMBED_DIM
    )
    path = tmp_path / "object_embeddings.npz"
    np.savez_compressed(path, embeddings=embeddings)
    metadata = {
        "task_id": "PickCube-v1",
        "instruction": "pick the cube",
        "object_embedding_path": f"{path.name}#embeddings/1",
        "_chart_root": str(tmp_path),
        "outcome_vector": {"success": False},
        "label": "negative",
    }

    feature = build_chart_feature(base_action, metadata, mode="base_context_obj")

    expected = base_action.size + 2 * CONTEXT_HASH_WIDTH + 2 + OBJECT_LAYOUT_EMBED_DIM
    assert feature.shape == (expected,)
    assert np.allclose(feature[-OBJECT_LAYOUT_EMBED_DIM:], embeddings[1])
    assert chart_feature_dim(base_action, mode="base_context_obj") == expected


def test_base_context_obs_obj_appends_both_embeddings(tmp_path) -> None:
    base_action = np.zeros((1, 4), dtype=np.float32)
    obs = np.ones((1, OBSERVATION_EMBED_DIM), dtype=np.float32)
    obj = np.full((1, OBJECT_LAYOUT_EMBED_DIM), 2.0, dtype=np.float32)
    obs_path = tmp_path / "obs.npz"
    obj_path = tmp_path / "obj.npz"
    np.savez_compressed(obs_path, embeddings=obs)
    np.savez_compressed(obj_path, embeddings=obj)
    metadata = {
        "task_id": "StackCube-v1",
        "instruction": "stack the cube",
        "observation_embedding_path": f"{obs_path.name}#embeddings/0",
        "object_embedding_path": f"{obj_path.name}#embeddings/0",
        "_chart_root": str(tmp_path),
    }

    feature = build_chart_feature(base_action, metadata, mode="base_context_obs_obj")

    expected = (
        base_action.size
        + 2 * CONTEXT_HASH_WIDTH
        + 2
        + OBSERVATION_EMBED_DIM
        + OBJECT_LAYOUT_EMBED_DIM
    )
    assert feature.shape == (expected,)
    assert np.allclose(feature[-(OBSERVATION_EMBED_DIM + OBJECT_LAYOUT_EMBED_DIM):-OBJECT_LAYOUT_EMBED_DIM], obs[0])
    assert np.allclose(feature[-OBJECT_LAYOUT_EMBED_DIM:], obj[0])


def test_unknown_chart_feature_mode_raises() -> None:
    with pytest.raises(ValueError, match="unknown chart feature mode"):
        build_chart_feature(np.zeros((1, 2), dtype=np.float32), mode="outcome")