tests/test_indra_engine.py

"""
Unit tests for the INDRA Poincaré Projection Engine (Sprint 29).

Tests cover:
1. Basic projection correctness (origin maps to origin)
2. Radial boundary enforcement (no vector exceeds R_MAX)
3. Floating-point variance underflow guard (identical vectors → distance 0)
4. Division-by-zero guard (boundary vectors → finite distance)
5. Sorting correctness (closer vectors rank higher)
6. Buffer reuse (no allocation leaks across calls)
"""

import numpy as np
import pytest
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from indra_engine import IndraProjectionEngine


@pytest.fixture
def engine():
    """Create engine with small batch for testing."""
    return IndraProjectionEngine(batch_size=10, dimensions=4)


class TestPoincareBallProjection:
    def test_zero_vector_maps_to_origin(self, engine):
        """A zero Euclidean vector should map to the Poincaré ball origin."""
        vec = np.zeros(4, dtype=np.float32)
        out = np.zeros(4, dtype=np.float32)
        engine.map_to_poincare_ball_inplace(vec, out)
        np.testing.assert_allclose(out, 0.0, atol=1e-7)

    def test_radial_boundary_clip(self, engine):
        """No projected vector should have norm > R_MAX."""
        vec = np.array([100.0, 200.0, 300.0, 400.0], dtype=np.float32)
        out = np.zeros(4, dtype=np.float32)
        engine.map_to_poincare_ball_inplace(vec, out)
        assert np.linalg.norm(out) <= engine.R_MAX + 1e-6

    def test_batch_radial_boundary_clip(self, engine):
        """No batch-projected vector should exceed R_MAX."""
        vecs = np.random.randn(5, 4).astype(np.float32) * 100
        out = np.zeros((10, 4), dtype=np.float32)
        engine.map_to_poincare_ball_inplace(vecs, out[:5])
        for i in range(5):
            assert np.linalg.norm(out[i]) <= engine.R_MAX + 1e-6


class TestPoincaréDistances:
    def test_identical_vectors_zero_distance(self, engine):
        """Identical vectors should produce distance ≈ 0."""
        query = np.array([0.5, 0.3, -0.2, 0.1], dtype=np.float32)
        candidates = np.tile(query, (3, 1))
        distances = engine.compute_poincare_distances(query, candidates, 3)
        np.testing.assert_allclose(distances, 0.0, atol=1e-5)

    def test_distance_ordering_preserved(self, engine):
        """Closer Euclidean vectors should also be closer in Poincaré space."""
        query = np.array([0.1, 0.1, 0.1, 0.1], dtype=np.float32)
        candidates = np.array([
            [0.11, 0.1, 0.1, 0.1],   # Very close
            [0.5, 0.5, 0.5, 0.5],     # Medium
            [2.0, 2.0, 2.0, 2.0],     # Far
        ], dtype=np.float32)
        distances = engine.compute_poincare_distances(query, candidates, 3)
        assert distances[0] < distances[1] < distances[2]

    def test_no_nan_or_inf(self, engine):
        """No NaN or Inf values in output, even with extreme inputs."""
        query = np.array([1e-10, 1e-10, 1e-10, 1e-10], dtype=np.float32)
        candidates = np.array([
            [0.0, 0.0, 0.0, 0.0],       # Zero vector
            [1e10, 1e10, 1e10, 1e10],     # Huge vector
            [1e-10, 1e-10, 1e-10, 1e-10], # Identical to query
        ], dtype=np.float32)
        distances = engine.compute_poincare_distances(query, candidates, 3)
        assert not np.any(np.isnan(distances))
        assert not np.any(np.isinf(distances))

    def test_boundary_vectors_finite(self, engine):
        """Vectors near the Poincaré ball boundary should produce finite distances."""
        query = np.array([0.998, 0.0, 0.0, 0.0], dtype=np.float32)
        candidates = np.array([
            [0.0, 0.998, 0.0, 0.0],
            [-0.998, 0.0, 0.0, 0.0],
        ], dtype=np.float32)
        distances = engine.compute_poincare_distances(query, candidates, 2)
        assert np.all(np.isfinite(distances))


class TestProjectAndRank:
    def test_returns_sorted_indices(self, engine):
        """project_and_rank should return indices sorted by ascending distance."""
        query = np.zeros(4, dtype=np.float32)
        candidates = np.array([
            [5.0, 5.0, 5.0, 5.0],     # Farthest
            [0.01, 0.01, 0.01, 0.01],  # Closest
            [1.0, 1.0, 1.0, 1.0],      # Middle
        ], dtype=np.float32)
        indices = engine.project_and_rank(query, candidates, 3)
        assert indices[0] == 1  # Closest first
        assert indices[-1] == 0  # Farthest last

    def test_buffer_reuse_stability(self, engine):
        """Running multiple times should produce consistent results (no buffer aliasing)."""
        query = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
        candidates = np.random.randn(5, 4).astype(np.float32)

        result1 = engine.project_and_rank(query, candidates, 5)
        result2 = engine.project_and_rank(query, candidates, 5)
        np.testing.assert_array_equal(result1, result2)