feox-ann 0.1.0

use feox_ann::{AnnConfig, AnnIndex, AnnQuery};

const DIMENSIONS: usize = 128;
const RECORDS: usize = 10_000;
const QUERIES: usize = 200;
const TOP_K: usize = 10;

struct XorShift(u64);

impl XorShift {
    fn next_f32(&mut self) -> f32 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        ((self.0 >> 40) as f32 / (1 << 24) as f32) * 2.0 - 1.0
    }

    fn vector(&mut self, dimensions: usize) -> Vec<f32> {
        (0..dimensions).map(|_| self.next_f32()).collect()
    }
}

fn normalize(vector: &mut [f32]) {
    let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
    if norm > 0.0 {
        for value in vector.iter_mut() {
            *value /= norm;
        }
    }
}

fn dot(left: &[f32], right: &[f32]) -> f32 {
    left.iter().zip(right).map(|(l, r)| l * r).sum()
}

fn main() -> feox_ann::Result<()> {
    let mut rng = XorShift(0x5eed_f0f5 ^ RECORDS as u64);
    let mut records: Vec<(String, Vec<f32>)> = (0..RECORDS)
        .map(|i| {
            let mut vector = rng.vector(DIMENSIONS);
            normalize(&mut vector);
            (format!("record-{i:06}"), vector)
        })
        .collect();

    let threads = std::thread::available_parallelism()
        .map(|threads| threads.get())
        .unwrap_or(1);
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(DIMENSIONS))?;
    index.bulk_load(
        records
            .iter()
            .map(|(id, vector)| (id.clone(), vector.clone())),
        threads,
    )?;

    let mut query_rng = XorShift(0xfeed_beef);
    let queries: Vec<Vec<f32>> = (0..QUERIES)
        .map(|_| {
            let mut vector = query_rng.vector(DIMENSIONS);
            normalize(&mut vector);
            vector
        })
        .collect();

    println!("{RECORDS} records x {DIMENSIONS}d, {QUERIES} queries, recall@{TOP_K}:");
    for ef_search in [32_usize, 64, 128, 256, 512] {
        let mut hits = 0_usize;
        for query in &queries {
            records.sort_by(|left, right| dot(query, &right.1).total_cmp(&dot(query, &left.1)));
            let truth: Vec<&str> = records[..TOP_K].iter().map(|(id, _)| id.as_str()).collect();

            let approx = index.query(AnnQuery {
                vector: query,
                top_k: TOP_K,
                ef_search: Some(ef_search),
                filter: None,
            })?;
            hits += approx
                .iter()
                .filter(|candidate| truth.contains(&candidate.id.as_str()))
                .count();
        }
        let recall = hits as f64 / (QUERIES * TOP_K) as f64;
        println!("  ef_search={ef_search:<4} recall@{TOP_K} = {recall:.4}");
    }

    Ok(())
}