brainy/tests/unit/hnsw/quantization.test.ts

/**
 * SQ8 Quantization Tests
 *
 * Tests for:
 * - SQ8 quantize/dequantize round-trip accuracy
 * - SQ8 distance vs float32 distance correlation
 * - Serialization/deserialization round-trip
 * - Search with quantization enabled: recall vs exact
 * - Two-phase rerank: improved recall over single-phase SQ8
 * - Config defaults: quantization disabled by default (no behavior change)
 */

import { describe, it, expect, beforeEach, afterEach } from 'vitest'
import { v4 as uuidv4 } from 'uuid'
import {
  quantizeSQ8,
  dequantizeSQ8,
  distanceSQ8,
  distanceSQ8Js,
  setSQ8DistanceImplementation,
  serializeSQ8,
  deserializeSQ8
} from '../../../src/utils/vectorQuantization.js'
import { JsHnswVectorIndex } from '../../../src/hnsw/hnswIndex.js'
import { euclideanDistance } from '../../../src/utils/index.js'
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'

// Helper: generate a random vector of given dimension
function randomVector(dim: number): number[] {
  return Array.from({ length: dim }, () => Math.random() * 2 - 1)
}

describe('SQ8 distance provider hook (setSQ8DistanceImplementation)', () => {
  const q1 = quantizeSQ8([0.1, 0.5, -0.3, 0.8])
  const q2 = quantizeSQ8([0.2, 0.4, -0.1, 0.7])

  // Always restore the JS default so a swapped impl never leaks into other tests.
  afterEach(() => setSQ8DistanceImplementation(distanceSQ8Js))

  it('dispatches to the JS implementation by default (identical result)', () => {
    const viaDispatcher = distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
    const viaJs = distanceSQ8Js(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
    expect(viaDispatcher).toBe(viaJs)
    expect(viaDispatcher).toBeGreaterThanOrEqual(0)
    expect(viaDispatcher).toBeLessThanOrEqual(2)
  })

  it('routes through a swapped implementation with the cortex 6-arg signature', () => {
    const calls: Array<[Uint8Array, number, number, Uint8Array, number, number]> = []
    setSQ8DistanceImplementation((a, aMin, aMax, b, bMin, bMax) => {
      calls.push([a, aMin, aMax, b, bMin, bMax])
      return 0.4242
    })
    const result = distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
    expect(result).toBe(0.4242)
    expect(calls).toHaveLength(1)
    // Exactly the arguments cortex's native cosineDistanceSq8(a, aMin, aMax, b, bMin, bMax) expects.
    expect(calls[0][0]).toBe(q1.quantized)
    expect(calls[0][1]).toBe(q1.min)
    expect(calls[0][3]).toBe(q2.quantized)
    expect(calls[0][5]).toBe(q2.max)
  })

  it('restores the JS default when set back', () => {
    setSQ8DistanceImplementation(() => -999)
    expect(distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)).toBe(-999)
    setSQ8DistanceImplementation(distanceSQ8Js)
    expect(distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)).toBe(
      distanceSQ8Js(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
    )
  })
})

// Helper: cosine distance for float32 vectors (reference implementation)
function cosineDistance(a: number[], b: number[]): number {
  let dot = 0
  let normA = 0
  let normB = 0
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i]
    normA += a[i] * a[i]
    normB += b[i] * b[i]
  }
  const denom = Math.sqrt(normA) * Math.sqrt(normB)
  if (denom === 0) return 1.0
  return 1 - dot / denom
}

describe('SQ8 Quantization', () => {
  // =================================================================
  // 1. QUANTIZE / DEQUANTIZE ROUND-TRIP
  // =================================================================
  describe('quantize/dequantize round-trip', () => {
    it('should round-trip a normalized vector with low error', () => {
      const original = randomVector(384)
      const sq8 = quantizeSQ8(original)
      const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max)

      expect(restored.length).toBe(original.length)

      // Per-dimension error should be small (max 1/255 of range)
      const range = sq8.max - sq8.min
      const maxError = range / 255
      for (let i = 0; i < original.length; i++) {
        expect(Math.abs(restored[i] - original[i])).toBeLessThanOrEqual(maxError + 1e-7)
      }
    })

    it('should handle a zero-range vector (all identical values)', () => {
      const original = new Array(128).fill(0.5)
      const sq8 = quantizeSQ8(original)

      expect(sq8.min).toBe(0.5)
      expect(sq8.max).toBe(0.5)
      // All quantized values should be 128 (midpoint)
      for (let i = 0; i < sq8.quantized.length; i++) {
        expect(sq8.quantized[i]).toBe(128)
      }

      const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max)
      for (let i = 0; i < restored.length; i++) {
        expect(restored[i]).toBe(0.5)
      }
    })

    it('should handle negative values', () => {
      const original = [-1.0, -0.5, 0, 0.5, 1.0]
      const sq8 = quantizeSQ8(original)

      expect(sq8.min).toBe(-1.0)
      expect(sq8.max).toBe(1.0)
      // Min maps to 0, max maps to 255
      expect(sq8.quantized[0]).toBe(0)   // -1.0
      expect(sq8.quantized[4]).toBe(255) // 1.0

      const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max)
      expect(Math.abs(restored[0] - (-1.0))).toBeLessThan(0.01)
      expect(Math.abs(restored[4] - 1.0)).toBeLessThan(0.01)
    })

    it('should produce Uint8Array output in [0, 255] range', () => {
      const original = randomVector(512)
      const sq8 = quantizeSQ8(original)

      expect(sq8.quantized).toBeInstanceOf(Uint8Array)
      expect(sq8.quantized.length).toBe(512)
      for (let i = 0; i < sq8.quantized.length; i++) {
        expect(sq8.quantized[i]).toBeGreaterThanOrEqual(0)
        expect(sq8.quantized[i]).toBeLessThanOrEqual(255)
      }
    })

    it('should achieve 4x storage reduction (uint8 vs float32)', () => {
      const dim = 384
      const original = randomVector(dim)
      const sq8 = quantizeSQ8(original)

      const float32Size = dim * 4 // 4 bytes per float32
      const sq8Size = sq8.quantized.byteLength + 8 // uint8 array + 2 floats for min/max
      const ratio = float32Size / sq8Size

      // Should be close to 4x (actually slightly less due to min/max overhead)
      expect(ratio).toBeGreaterThan(3.9)
    })
  })

  // =================================================================
  // 2. SQ8 DISTANCE VS FLOAT32 DISTANCE
  // =================================================================
  describe('SQ8 distance accuracy', () => {
    it('should correlate strongly with float32 cosine distance', () => {
      const dim = 384
      const numPairs = 100
      const errors: number[] = []

      for (let i = 0; i < numPairs; i++) {
        const a = randomVector(dim)
        const b = randomVector(dim)

        const exactDist = cosineDistance(a, b)

        const sq8A = quantizeSQ8(a)
        const sq8B = quantizeSQ8(b)
        const approxDist = distanceSQ8(
          sq8A.quantized, sq8A.min, sq8A.max,
          sq8B.quantized, sq8B.min, sq8B.max
        )

        errors.push(Math.abs(exactDist - approxDist))
      }

      // Mean absolute error should be very small
      const meanError = errors.reduce((sum, e) => sum + e, 0) / errors.length
      expect(meanError).toBeLessThan(0.02) // Less than 2% average error

      // Max error should be bounded
      const maxError = Math.max(...errors)
      expect(maxError).toBeLessThan(0.1) // Less than 10% worst case
    })

    it('should return 0 for identical vectors', () => {
      const v = randomVector(128)
      const sq8 = quantizeSQ8(v)
      const dist = distanceSQ8(
        sq8.quantized, sq8.min, sq8.max,
        sq8.quantized, sq8.min, sq8.max
      )
      expect(dist).toBeCloseTo(0, 5)
    })

    it('should preserve relative ordering of distances', () => {
      const dim = 128
      const query = randomVector(dim)
      // Create a vector close to query and one far away
      const close = query.map(v => v + (Math.random() * 0.1 - 0.05))
      const far = randomVector(dim)

      const exactClose = cosineDistance(query, close)
      const exactFar = cosineDistance(query, far)

      const sq8Query = quantizeSQ8(query)
      const sq8Close = quantizeSQ8(close)
      const sq8Far = quantizeSQ8(far)

      const approxClose = distanceSQ8(
        sq8Query.quantized, sq8Query.min, sq8Query.max,
        sq8Close.quantized, sq8Close.min, sq8Close.max
      )
      const approxFar = distanceSQ8(
        sq8Query.quantized, sq8Query.min, sq8Query.max,
        sq8Far.quantized, sq8Far.min, sq8Far.max
      )

      // Relative ordering should be preserved
      if (exactClose < exactFar) {
        expect(approxClose).toBeLessThan(approxFar)
      }
    })

    it('should handle zero vectors gracefully', () => {
      const zero = new Array(64).fill(0)
      const nonZero = randomVector(64)

      const sq8Zero = quantizeSQ8(zero)
      const sq8NonZero = quantizeSQ8(nonZero)

      const dist = distanceSQ8(
        sq8Zero.quantized, sq8Zero.min, sq8Zero.max,
        sq8NonZero.quantized, sq8NonZero.min, sq8NonZero.max
      )
      // Zero vectors should return max distance (1.0)
      expect(dist).toBeCloseTo(1.0, 1)
    })
  })

  // =================================================================
  // 3. SERIALIZATION / DESERIALIZATION
  // =================================================================
  describe('serialization round-trip', () => {
    it('should serialize and deserialize SQ8 data with float32 precision', () => {
      const original = randomVector(384)
      const sq8 = quantizeSQ8(original)
      const buffer = serializeSQ8(sq8)
      const restored = deserializeSQ8(buffer)

      // min/max are stored as float32, so some precision loss is expected
      expect(restored.min).toBeCloseTo(sq8.min, 5)
      expect(restored.max).toBeCloseTo(sq8.max, 5)
      expect(restored.quantized.length).toBe(sq8.quantized.length)
      for (let i = 0; i < sq8.quantized.length; i++) {
        expect(restored.quantized[i]).toBe(sq8.quantized[i])
      }
    })

    it('should produce compact binary format (8 + dim bytes)', () => {
      const dim = 384
      const sq8 = quantizeSQ8(randomVector(dim))
      const buffer = serializeSQ8(sq8)

      // 4 bytes for min (float32) + 4 bytes for max (float32) + dim bytes
      expect(buffer.byteLength).toBe(8 + dim)
    })
  })

  // =================================================================
  // 4. HNSW SEARCH WITH QUANTIZATION
  // =================================================================
  describe('HNSW search with quantization', () => {
    let index: JsHnswVectorIndex
    let storage: MemoryStorage
    const dim = 32 // Small dimension for fast tests
    let ids: string[]

    beforeEach(async () => {
      storage = new MemoryStorage()
      index = new JsHnswVectorIndex(
        {
          M: 8,
          efConstruction: 100,
          efSearch: 50,
          ml: 8,
          quantization: { enabled: true, rerankMultiplier: 3 }
        },
        euclideanDistance,
        { useParallelization: false, storage }
      )
      ids = []
    })

    it('should return correct results with quantization enabled', async () => {
      // Insert entities with UUID IDs
      const vectors: number[][] = []
      for (let i = 0; i < 50; i++) {
        const id = uuidv4()
        ids.push(id)
        const v = randomVector(dim)
        vectors.push(v)
        await index.addItem({ id, vector: v })
      }

      // Search with a known vector
      const queryIdx = 5
      const results = await index.search(vectors[queryIdx], 5)

      // The exact vector should be the closest (distance ~0)
      expect(results.length).toBeGreaterThan(0)
      expect(results[0][0]).toBe(ids[queryIdx])
      expect(results[0][1]).toBeCloseTo(0, 1)
    })

    it('should find the exact match as top result', async () => {
      const targetId = uuidv4()
      const target = randomVector(dim)
      await index.addItem({ id: targetId, vector: target })

      // Add noise vectors
      for (let i = 0; i < 30; i++) {
        await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
      }

      const results = await index.search(target, 1)
      expect(results.length).toBe(1)
      expect(results[0][0]).toBe(targetId)
    })

    it('should respect k parameter', async () => {
      for (let i = 0; i < 100; i++) {
        await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
      }

      const results5 = await index.search(randomVector(dim), 5)
      const results10 = await index.search(randomVector(dim), 10)

      expect(results5.length).toBe(5)
      // With quantization on a small graph (100 items), HNSW may occasionally
      // return slightly fewer than k due to approximation in distance calculations
      expect(results10.length).toBeGreaterThanOrEqual(8)
      expect(results10.length).toBeLessThanOrEqual(10)
    })
  })

  // =================================================================
  // 5. TWO-PHASE RERANK
  // =================================================================
  describe('two-phase rerank', () => {
    it('should accept rerank options in search', async () => {
      const storage = new MemoryStorage()
      const index = new JsHnswVectorIndex(
        {
          M: 8,
          efConstruction: 100,
          efSearch: 50,
          ml: 8,
          quantization: { enabled: true, rerankMultiplier: 1 } // Disable default rerank
        },
        euclideanDistance,
        { useParallelization: false, storage }
      )

      const dim = 32
      for (let i = 0; i < 30; i++) {
        await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
      }

      // Search with explicit rerank
      const results = await index.search(
        randomVector(dim),
        5,
        undefined,
        { rerank: { multiplier: 3 } }
      )

      expect(results.length).toBe(5)
      // Results should be sorted by distance
      for (let i = 1; i < results.length; i++) {
        expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
      }
    })

    it('reranked results should be sorted by exact distance', async () => {
      const storage = new MemoryStorage()
      const index = new JsHnswVectorIndex(
        {
          M: 8,
          efConstruction: 100,
          efSearch: 50,
          ml: 8,
          quantization: { enabled: true, rerankMultiplier: 3 }
        },
        euclideanDistance,
        { useParallelization: false, storage }
      )

      const dim = 32
      for (let i = 0; i < 50; i++) {
        await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
      }

      const query = randomVector(dim)
      const results = await index.search(query, 10)

      // Results should be sorted by exact distance (after reranking)
      for (let i = 1; i < results.length; i++) {
        expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
      }
    })
  })

  // =================================================================
  // 6. CONFIG DEFAULTS
  // =================================================================
  describe('configuration defaults', () => {
    it('should disable quantization by default', async () => {
      const storage = new MemoryStorage()
      const defaultIndex = new JsHnswVectorIndex(
        { M: 4, efConstruction: 50, efSearch: 20 },
        euclideanDistance,
        { useParallelization: false, storage }
      )

      const dim = 16
      for (let i = 0; i < 10; i++) {
        await defaultIndex.addItem({ id: uuidv4(), vector: randomVector(dim) })
      }

      // Search should work identically to pre-quantization behavior
      const results = await defaultIndex.search(randomVector(dim), 5)
      expect(results.length).toBe(5)

      // Distances should be exact euclidean (not SQ8 approximate)
      for (const [, dist] of results) {
        expect(typeof dist).toBe('number')
        expect(dist).toBeGreaterThanOrEqual(0)
      }
    })

    it('should use default rerankMultiplier of 3', async () => {
      const storage = new MemoryStorage()
      const index = new JsHnswVectorIndex(
        {
          M: 4,
          efConstruction: 50,
          efSearch: 20,
          quantization: { enabled: true }
        },
        euclideanDistance,
        { useParallelization: false, storage }
      )

      // The index should be created without errors
      // Default rerankMultiplier should be 3 (tested implicitly via search)
      const dim = 16
      for (let i = 0; i < 10; i++) {
        await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
      }

      const results = await index.search(randomVector(dim), 3)
      expect(results.length).toBe(3)
    })

    it('should default vectorStorage to memory mode', async () => {
      const storage = new MemoryStorage()
      const index = new JsHnswVectorIndex(
        { M: 4, efConstruction: 50, efSearch: 20 },
        euclideanDistance,
        { useParallelization: false, storage }
      )

      const id = uuidv4()
      const v = randomVector(16)
      await index.addItem({ id, vector: v })

      // In memory mode, vector should be immediately searchable
      const results = await index.search(v, 1)
      expect(results.length).toBe(1)
      expect(results[0][0]).toBe(id)
      expect(results[0][1]).toBeCloseTo(0, 5)
    })
  })
})