evoc-rs 0.3.0

#![cfg(feature = "gpu")]
#![allow(clippy::needless_range_loop)]

mod commons;
use commons::*;

use ann_search_rs::synthetic::generate_clustered_data;
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};
use faer::Mat;

use evoc_rs::nearest_neighbours::nearest_neighbour_gpu::run_ann_search_gpu;
use evoc_rs::prelude::*;
use evoc_rs::{EvocParams, evoc_gpu};

/// Fixture for the end-to-end accuracy tests: four clusters plus inter-cluster
/// bridges. Over 30 EVoC seeds this lands in [0.965, 0.998] on every CPU and
/// GPU backend, whereas `make_blobs` swings across the whole [0.49, 0.99] range
/// with the data seed.
const ACC_FIXTURE: (usize, usize, usize, u64) = (600, 32, 4, 1);

/// Convert `Vec<Vec<f64>>` from make_blobs into a faer matrix.
fn to_mat(data: &[Vec<f64>]) -> Mat<f32> {
    let n = data.len();
    let d = data[0].len();
    Mat::from_fn(n, d, |i, j| data[i][j] as f32)
}

/// End-to-end GPU EVoC on clustered data should recover the clusters.
#[test]
fn gpu_integration_01_clusters_exhaustive() {
    let (n, dim, nc, seed) = ACC_FIXTURE;
    let (mat, gt) = generate_clustered_data::<f32>(n, dim, nc, seed);

    let params = EvocParams::<f32>::default();
    let nn_params = NearestNeighbourParamsGpuEvoc::<f32>::default();
    let device = WgpuDevice::default();

    let result = evoc_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        "exhaustive_gpu".to_string(),
        None,
        &params,
        &nn_params,
        device,
        42,
        0,
    )
    .unwrap();

    assert!(!result.cluster_layers.is_empty());
    let acc = cluster_accuracy(result.best_labels(), &gt);
    println!(
        "Accuracy (exhaustive_gpu): {:.3}, k = {}",
        acc,
        result.n_clusters()
    );
    assert!(acc > 0.95, "Accuracy {:.3} below threshold", acc);
}

/// Same test with ivf_gpu, less precise but should still separate clusters.
#[test]
fn gpu_integration_02_clusters_ivf() {
    let (n, dim, nc, seed) = ACC_FIXTURE;
    let (mat, gt) = generate_clustered_data::<f32>(n, dim, nc, seed);

    let params = EvocParams::<f32>::default();
    let nn_params = NearestNeighbourParamsGpuEvoc::<f32> {
        n_list: Some(5),
        n_probes: Some(5),
        ..NearestNeighbourParamsGpuEvoc::default()
    };

    let device = WgpuDevice::default();

    let result = evoc_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        "ivf_gpu".to_string(),
        None,
        &params,
        &nn_params,
        device,
        42,
        0,
    )
    .unwrap();

    let acc = cluster_accuracy(result.best_labels(), &gt);
    println!(
        "Accuracy (ivf_gpu): {:.3}, k = {}",
        acc,
        result.n_clusters()
    );
    assert!(acc > 0.95, "IVF accuracy {:.3} too low", acc);
}

/// Dispatch check — all three GPU backends should run without panicking.
#[test]
fn gpu_integration_03_all_backends_dispatch() {
    let (data, _) = make_blobs(100, 3, 8, 20.0, 5.0, 42);
    let mat = to_mat(&data);

    let params = EvocParams::<f32>::default();
    let nn_params = NearestNeighbourParamsGpuEvoc::<f32>::default();

    for ann_type in ["exhaustive_gpu", "ivf_gpu", "nndescent_gpu"] {
        let device = WgpuDevice::default();
        let result = evoc_gpu::<f32, WgpuRuntime>(
            mat.as_ref(),
            ann_type.to_string(),
            None,
            &params,
            &nn_params,
            device,
            42,
            0,
        )
        .unwrap();
        assert!(
            !result.cluster_layers.is_empty(),
            "{} produced no layers",
            ann_type
        );
        assert_eq!(result.nn_indices.len(), data.len());
        println!("{} dispatched ok", ann_type);
    }
}

/// GPU result should be structurally consistent with CPU result on the same
/// data (not identical — GPU isn't bit-reproducible — but similar cluster count
/// and comparable accuracy).
#[test]
fn gpu_integration_04_structural_agreement_with_cpu() {
    use evoc_rs::evoc;

    let (n, dim, nc, seed) = ACC_FIXTURE;
    let (mat, gt) = generate_clustered_data::<f32>(n, dim, nc, seed);

    let params = EvocParams::<f32>::default();

    let nn_cpu = NearestNeighbourParamsEvoc::<f32>::default();
    let cpu = evoc::<f32>(
        mat.as_ref(),
        "kmknn".to_string(),
        None,
        &params,
        &nn_cpu,
        42,
        0,
    )
    .unwrap();

    let nn_gpu = NearestNeighbourParamsGpuEvoc::<f32>::default();
    let device = WgpuDevice::default();
    let gpu = evoc_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        "exhaustive_gpu".to_string(),
        None,
        &params,
        &nn_gpu,
        device,
        42,
        0,
    )
    .unwrap();

    let acc_cpu = cluster_accuracy(cpu.best_labels(), &gt);
    let acc_gpu = cluster_accuracy(gpu.best_labels(), &gt);
    println!("CPU accuracy: {:.3}, GPU accuracy: {:.3}", acc_cpu, acc_gpu);

    assert!(acc_cpu > 0.9 && acc_gpu > 0.9);
    let k_cpu = cpu.n_clusters();
    let k_gpu = gpu.n_clusters();
    println!("Cluster counts - CPU: {}, GPU: {}", k_cpu, k_gpu);
    let diff = (k_cpu as isize - k_gpu as isize).unsigned_abs();
    assert!(
        diff <= 2,
        "CPU/GPU cluster counts diverge: {} vs {}",
        k_cpu,
        k_gpu
    );
}

/// Precomputed kNN path should produce a valid clustering.
#[test]
fn gpu_integration_05_precomputed_knn() {
    let (n, dim, nc, seed) = ACC_FIXTURE;
    let (mat, gt) = generate_clustered_data::<f32>(n, dim, nc, seed);
    let k = 15;

    let nn_params = NearestNeighbourParamsGpuEvoc::<f32>::default();
    let device = WgpuDevice::default();
    let (knn_indices, knn_dist) = run_ann_search_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        k,
        "exhaustive_gpu".to_string(),
        &nn_params,
        device,
        42,
        0,
    )
    .unwrap();

    let params = EvocParams::<f32> {
        n_neighbours: k,
        ..EvocParams::default()
    };

    let device2 = WgpuDevice::default();
    let result = evoc_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        "exhaustive_gpu".to_string(),
        Some((knn_indices, knn_dist)),
        &params,
        &nn_params,
        device2,
        42,
        0,
    )
    .unwrap();

    let acc = cluster_accuracy(result.best_labels(), &gt);
    assert!(acc > 0.95, "Precomputed-kNN accuracy {:.3} too low", acc);
}

/// approx_n_clusters mode with GPU kNN.
#[test]
fn gpu_integration_06_approx_n_clusters() {
    let (data, _) = make_blobs(100, 4, 4, 20.0, 0.5, 42);
    let mat = to_mat(&data);

    let params = EvocParams::<f32> {
        approx_n_clusters: Some(4),
        min_samples: 3,
        base_min_cluster_size: 3,
        ..EvocParams::default()
    };
    let nn_params = NearestNeighbourParamsGpuEvoc::<f32>::default();
    let device = WgpuDevice::default();

    let result = evoc_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        "exhaustive_gpu".to_string(),
        None,
        &params,
        &nn_params,
        device,
        42,
        0,
    )
    .unwrap();

    assert_eq!(
        result.cluster_layers.len(),
        1,
        "approx_n_clusters should yield one layer"
    );
    let k = result.n_clusters();
    println!("Requested 4 clusters, got {}", k);
    assert!((2..=6).contains(&k), "Got {} clusters, expected near 4", k);
}

/// kNN indices returned from GPU path exclude self (latent run_ann_search_gpu
/// concern noted elsewhere — verify here at the EVoC boundary).
#[test]
fn gpu_integration_07_knn_no_self() {
    let (data, _) = make_blobs(100, 2, 4, 20.0, 5.0, 42);
    let mat = to_mat(&data);

    let params = EvocParams::<f32>::default();
    let nn_params = NearestNeighbourParamsGpuEvoc::<f32>::default();
    let device = WgpuDevice::default();

    let result = evoc_gpu::<f32, WgpuRuntime>(
        mat.as_ref(),
        "exhaustive_gpu".to_string(),
        None,
        &params,
        &nn_params,
        device,
        42,
        0,
    )
    .unwrap();

    let mut self_loops = 0;
    for (i, nb) in result.nn_indices.iter().enumerate() {
        if nb.contains(&i) {
            self_loops += 1;
        }
    }
    // Accept up to 1% as acknowledged GPU tie-breaking artefact.
    let max_allowed = result.nn_indices.len() / 100 + 1;
    assert!(
        self_loops <= max_allowed,
        "{} self-loops out of {} (>1%)",
        self_loops,
        result.nn_indices.len()
    );
}