diff --git a/mini-nav/scenegraph/__init__.py b/mini-nav/scenegraph/__init__.py index a94d20e..5adec01 100644 --- a/mini-nav/scenegraph/__init__.py +++ b/mini-nav/scenegraph/__init__.py @@ -5,8 +5,24 @@ This module exports the main scenegraph objects for easy import: from mini_nav.scenegraph import SimpleSceneGraph, RoomNode, ObjectNode """ +from .hash_codec import ( + bits_tensor_to_cam_row, + bits_tensor_to_hash_bytes, + cam_row_to_hash_bytes, + hash_bytes_to_bits_array, + hash_bytes_to_cam_row, +) from .objectnode import ObjectNode from .roomnode import RoomNode from .scenegraph import SimpleSceneGraph -__all__ = ["ObjectNode", "RoomNode", "SimpleSceneGraph"] +__all__ = [ + "ObjectNode", + "RoomNode", + "SimpleSceneGraph", + "bits_tensor_to_cam_row", + "bits_tensor_to_hash_bytes", + "cam_row_to_hash_bytes", + "hash_bytes_to_bits_array", + "hash_bytes_to_cam_row", +] diff --git a/mini-nav/scenegraph/hash_codec.py b/mini-nav/scenegraph/hash_codec.py new file mode 100644 index 0000000..98ce8b9 --- /dev/null +++ b/mini-nav/scenegraph/hash_codec.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import numpy as np +import torch + +DEFAULT_HASH_WIDTH = 512 + + +def _validate_width(width: int) -> None: + if width <= 0: + raise ValueError("width must be greater than 0") + if width % 8 != 0: + raise ValueError("width must be divisible by 8") + + +def _expected_byte_length(width: int) -> int: + _validate_width(width) + return width // 8 + + +def bits_tensor_to_hash_bytes(bits: torch.Tensor, *, width: int = 512) -> bytes: + _validate_width(width) + flat = bits.detach().cpu().flatten() + if flat.numel() != width: + raise ValueError( + f"Input tensor must have exactly {width} values, got {flat.numel()}" + ) + bit_array = (flat.numpy() > 0).astype(np.uint8) + return np.packbits(bit_array).tobytes() + + +def hash_bytes_to_bits_array(hash_bytes: bytes, *, width: int = 512) -> np.ndarray: + expected = _expected_byte_length(width) + if len(hash_bytes) != expected: + raise ValueError( + f"hash_bytes must be exactly {expected} bytes, got {len(hash_bytes)}" + ) + byte_array = np.frombuffer(hash_bytes, dtype=np.uint8) + return np.unpackbits(byte_array)[:width].astype(np.uint8, copy=False) + + +def hash_bytes_to_cam_row(hash_bytes: bytes, *, width: int = 512) -> int: + expected = _expected_byte_length(width) + if len(hash_bytes) != expected: + raise ValueError( + f"hash_bytes must be exactly {expected} bytes, got {len(hash_bytes)}" + ) + return int.from_bytes(hash_bytes, byteorder="big", signed=False) + + +def cam_row_to_hash_bytes(cam_row: int, *, width: int = 512) -> bytes: + if not (0 <= cam_row < 1 << width): + raise ValueError( + f"cam_row {cam_row} is out of range [0, 2**{width})" + ) + expected = _expected_byte_length(width) + return int(cam_row).to_bytes(expected, byteorder="big", signed=False) + + +def bits_tensor_to_cam_row(bits: torch.Tensor, *, width: int = 512) -> int: + hash_bytes = bits_tensor_to_hash_bytes(bits, width=width) + return hash_bytes_to_cam_row(hash_bytes, width=width) diff --git a/tests/test_hash_codec.py b/tests/test_hash_codec.py new file mode 100644 index 0000000..924ab47 --- /dev/null +++ b/tests/test_hash_codec.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + + +MINI_NAV_DIR = Path(__file__).resolve().parents[1] / "mini-nav" +sys.path.insert(0, str(MINI_NAV_DIR)) + +from scenegraph.hash_codec import ( # noqa: E402 + bits_tensor_to_cam_row, + bits_tensor_to_hash_bytes, + cam_row_to_hash_bytes, + hash_bytes_to_bits_array, + hash_bytes_to_cam_row, +) + + +WIDTH = 512 + + +def _xnor_score(query_row: int, stored_row: int, *, width: int = WIDTH) -> int: + mask = (1 << width) - 1 + return int((~(query_row ^ stored_row) & mask).bit_count()) + + +def _hamming_distance(left: np.ndarray, right: np.ndarray) -> int: + return int((left != right).sum()) + + +def test_all_zero_hash_roundtrips_through_bytes_and_cam_row(): + bits = torch.zeros(WIDTH, dtype=torch.int32) + + hash_bytes = bits_tensor_to_hash_bytes(bits) + cam_row = hash_bytes_to_cam_row(hash_bytes) + roundtrip = cam_row_to_hash_bytes(cam_row) + + assert len(hash_bytes) == WIDTH // 8 + assert hash_bytes == bytes(WIDTH // 8) + assert cam_row == 0 + assert roundtrip == hash_bytes + + +def test_all_one_hash_roundtrips_through_bytes_and_cam_row(): + bits = torch.ones(WIDTH, dtype=torch.int32) + + hash_bytes = bits_tensor_to_hash_bytes(bits) + cam_row = hash_bytes_to_cam_row(hash_bytes) + roundtrip = cam_row_to_hash_bytes(cam_row) + + assert hash_bytes == b"\xff" * (WIDTH // 8) + assert cam_row == (1 << WIDTH) - 1 + assert roundtrip == hash_bytes + + +def test_first_bit_uses_packbits_high_bit_ordering(): + bits = torch.zeros(WIDTH, dtype=torch.int32) + bits[0] = 1 + + hash_bytes = bits_tensor_to_hash_bytes(bits) + cam_row = hash_bytes_to_cam_row(hash_bytes) + + assert hash_bytes[0] == 0b10000000 + assert cam_row == 1 << (WIDTH - 1) + + +def test_last_bit_maps_to_least_significant_cam_row_bit(): + bits = torch.zeros(WIDTH, dtype=torch.int32) + bits[WIDTH - 1] = 1 + + hash_bytes = bits_tensor_to_hash_bytes(bits) + cam_row = hash_bytes_to_cam_row(hash_bytes) + + assert hash_bytes[-1] == 0b00000001 + assert cam_row == 1 + + +def test_hash_bytes_unpack_to_bits_array_with_packbits_ordering(): + hash_bytes = bytes([0b10100000]) + bytes((WIDTH // 8) - 1) + + bits = hash_bytes_to_bits_array(hash_bytes) + + assert bits.shape == (WIDTH,) + assert bits.dtype == np.uint8 + assert bits[:4].tolist() == [1, 0, 1, 0] + assert bits[4:].sum() == 0 + + +def test_bits_tensor_to_cam_row_matches_bytes_conversion(): + bits = torch.zeros(WIDTH, dtype=torch.int32) + bits[0] = 1 + bits[7] = 1 + bits[511] = 1 + + hash_bytes = bits_tensor_to_hash_bytes(bits) + + assert bits_tensor_to_cam_row(bits) == hash_bytes_to_cam_row(hash_bytes) + + +def test_positive_threshold_accepts_bool_and_signed_hash_encodings(): + bool_bits = torch.zeros(WIDTH, dtype=torch.bool) + bool_bits[0] = True + signed_bits = torch.full((WIDTH,), -1, dtype=torch.int32) + signed_bits[0] = 1 + + assert bits_tensor_to_hash_bytes(bool_bits) == bits_tensor_to_hash_bytes(signed_bits) + + +def test_xnor_score_matches_width_minus_hamming_distance(): + zeros = torch.zeros(WIDTH, dtype=torch.int32) + ones = torch.ones(WIDTH, dtype=torch.int32) + one_bit_diff = torch.zeros(WIDTH, dtype=torch.int32) + one_bit_diff[0] = 1 + + zero_bytes = bits_tensor_to_hash_bytes(zeros) + one_bytes = bits_tensor_to_hash_bytes(ones) + one_diff_bytes = bits_tensor_to_hash_bytes(one_bit_diff) + + zero_row = hash_bytes_to_cam_row(zero_bytes) + one_row = hash_bytes_to_cam_row(one_bytes) + one_diff_row = hash_bytes_to_cam_row(one_diff_bytes) + + zero_bits = hash_bytes_to_bits_array(zero_bytes) + one_bits = hash_bytes_to_bits_array(one_bytes) + one_diff_bits = hash_bytes_to_bits_array(one_diff_bytes) + + assert _hamming_distance(zero_bits, zero_bits) == 0 + assert _xnor_score(zero_row, zero_row) == WIDTH + + assert _hamming_distance(zero_bits, one_diff_bits) == 1 + assert _xnor_score(zero_row, one_diff_row) == WIDTH - 1 + + assert _hamming_distance(zero_bits, one_bits) == WIDTH + assert _xnor_score(zero_row, one_row) == 0 + + +@pytest.mark.parametrize( + "bits", + [ + torch.zeros(WIDTH - 1, dtype=torch.int32), + torch.zeros(WIDTH + 1, dtype=torch.int32), + ], +) +def test_bits_tensor_to_hash_bytes_rejects_wrong_width(bits: torch.Tensor): + with pytest.raises(ValueError, match="exactly 512 values"): + bits_tensor_to_hash_bytes(bits) + + +def test_hash_bytes_to_bits_array_rejects_wrong_byte_length(): + with pytest.raises(ValueError, match="exactly 64 bytes"): + hash_bytes_to_bits_array(bytes(63)) + + +@pytest.mark.parametrize("cam_row", [-1, 1 << WIDTH]) +def test_cam_row_to_hash_bytes_rejects_out_of_range_rows(cam_row: int): + with pytest.raises(ValueError, match="range"): + cam_row_to_hash_bytes(cam_row) + + +def test_width_must_be_divisible_by_eight(): + with pytest.raises(ValueError, match="divisible by 8"): + bits_tensor_to_hash_bytes(torch.zeros(7, dtype=torch.int32), width=7) + + +def test_scenegraph_package_exports_hash_codec_helpers(): + from scenegraph import ( # noqa: PLC0415 + bits_tensor_to_cam_row as exported_bits_tensor_to_cam_row, + bits_tensor_to_hash_bytes as exported_bits_tensor_to_hash_bytes, + cam_row_to_hash_bytes as exported_cam_row_to_hash_bytes, + hash_bytes_to_bits_array as exported_hash_bytes_to_bits_array, + hash_bytes_to_cam_row as exported_hash_bytes_to_cam_row, + ) + + assert exported_bits_tensor_to_cam_row is bits_tensor_to_cam_row + assert exported_bits_tensor_to_hash_bytes is bits_tensor_to_hash_bytes + assert exported_cam_row_to_hash_bytes is cam_row_to_hash_bytes + assert exported_hash_bytes_to_bits_array is hash_bytes_to_bits_array + assert exported_hash_bytes_to_cam_row is hash_bytes_to_cam_row