Skip to main content

scirs2_cluster/simd_ops/
distance.rs

1//! SIMD-accelerated distance computation for clustering algorithms.
2//!
3//! Provides Euclidean, Manhattan, pairwise, and condensed distance computation
4//! using SIMD acceleration from `scirs2-core`.
5
6use scirs2_core::ndarray::{s, Array1, Array2, ArrayView1, ArrayView2, Axis};
7use scirs2_core::numeric::{Float, FromPrimitive};
8use scirs2_core::parallel_ops::*;
9use scirs2_core::simd_ops::{AutoOptimizer, PlatformCapabilities, SimdUnifiedOps};
10use std::fmt::Debug;
11
12use crate::error::{ClusteringError, Result};
13
14// ---------------------------------------------------------------------------
15// Configuration
16// ---------------------------------------------------------------------------
17
18/// Configuration for SIMD-accelerated clustering operations.
19///
20/// Controls thresholds for enabling SIMD, parallel processing, and
21/// cache-friendly memory access patterns.
22#[derive(Debug, Clone)]
23pub struct SimdClusterConfig {
24    /// Minimum vector length to trigger SIMD optimizations (default: 32)
25    pub simd_threshold: usize,
26    /// Enable parallel processing for large workloads (default: true)
27    pub enable_parallel: bool,
28    /// Chunk size for parallel processing (default: 512)
29    pub parallel_chunk_size: usize,
30    /// Block size for cache-friendly pairwise distance computation (default: 128)
31    pub block_size: usize,
32    /// Force SIMD usage regardless of heuristics (for testing; default: false)
33    pub force_simd: bool,
34}
35
36impl Default for SimdClusterConfig {
37    fn default() -> Self {
38        Self {
39            simd_threshold: 32,
40            enable_parallel: true,
41            parallel_chunk_size: 512,
42            block_size: 128,
43            force_simd: false,
44        }
45    }
46}
47
48/// Determine whether SIMD should be used for a given problem size.
49pub(super) fn should_use_simd(n_elements: usize, config: &SimdClusterConfig) -> bool {
50    let caps = PlatformCapabilities::detect();
51    let optimizer = AutoOptimizer::new();
52    config.force_simd
53        || (caps.simd_available
54            && (optimizer.should_use_simd(n_elements) || n_elements >= config.simd_threshold))
55}
56
57// ---------------------------------------------------------------------------
58// Distance Computation
59// ---------------------------------------------------------------------------
60
61/// Distance metric for SIMD-accelerated operations.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum SimdDistanceMetric {
64    /// Euclidean distance (L2 norm)
65    Euclidean,
66    /// Squared Euclidean distance (avoids sqrt for speed)
67    SquaredEuclidean,
68    /// Manhattan distance (L1 norm)
69    Manhattan,
70}
71
72/// Compute Euclidean distance between two vectors using SIMD.
73///
74/// Falls back to scalar computation when SIMD is not beneficial.
75///
76/// # Errors
77///
78/// Returns `ClusteringError::InvalidInput` if vectors have different lengths.
79pub fn simd_euclidean_distance<F>(
80    a: ArrayView1<F>,
81    b: ArrayView1<F>,
82    config: Option<&SimdClusterConfig>,
83) -> Result<F>
84where
85    F: Float + FromPrimitive + Debug + SimdUnifiedOps,
86{
87    if a.len() != b.len() {
88        return Err(ClusteringError::InvalidInput(format!(
89            "Vectors must have the same length: got {} and {}",
90            a.len(),
91            b.len()
92        )));
93    }
94    if a.is_empty() {
95        return Ok(F::zero());
96    }
97    let default_config = SimdClusterConfig::default();
98    let cfg = config.unwrap_or(&default_config);
99
100    if should_use_simd(a.len(), cfg) {
101        let diff = F::simd_sub(&a, &b);
102        Ok(F::simd_norm(&diff.view()))
103    } else {
104        let mut sum = F::zero();
105        for i in 0..a.len() {
106            let d = a[i] - b[i];
107            sum = sum + d * d;
108        }
109        Ok(sum.sqrt())
110    }
111}
112
113/// Compute squared Euclidean distance between two vectors using SIMD.
114///
115/// This avoids the square root and is preferred when only relative ordering
116/// of distances matters (e.g., nearest-neighbor search in k-means).
117///
118/// # Errors
119///
120/// Returns `ClusteringError::InvalidInput` if vectors have different lengths.
121pub fn simd_squared_euclidean_distance<F>(
122    a: ArrayView1<F>,
123    b: ArrayView1<F>,
124    config: Option<&SimdClusterConfig>,
125) -> Result<F>
126where
127    F: Float + FromPrimitive + Debug + SimdUnifiedOps,
128{
129    if a.len() != b.len() {
130        return Err(ClusteringError::InvalidInput(format!(
131            "Vectors must have the same length: got {} and {}",
132            a.len(),
133            b.len()
134        )));
135    }
136    if a.is_empty() {
137        return Ok(F::zero());
138    }
139    let default_config = SimdClusterConfig::default();
140    let cfg = config.unwrap_or(&default_config);
141
142    if should_use_simd(a.len(), cfg) {
143        let diff = F::simd_sub(&a, &b);
144        let sq = F::simd_mul(&diff.view(), &diff.view());
145        Ok(F::simd_sum(&sq.view()))
146    } else {
147        let mut sum = F::zero();
148        for i in 0..a.len() {
149            let d = a[i] - b[i];
150            sum = sum + d * d;
151        }
152        Ok(sum)
153    }
154}
155
156/// Compute Manhattan distance between two vectors using SIMD.
157///
158/// # Errors
159///
160/// Returns `ClusteringError::InvalidInput` if vectors have different lengths.
161pub fn simd_manhattan_distance<F>(
162    a: ArrayView1<F>,
163    b: ArrayView1<F>,
164    config: Option<&SimdClusterConfig>,
165) -> Result<F>
166where
167    F: Float + FromPrimitive + Debug + SimdUnifiedOps,
168{
169    if a.len() != b.len() {
170        return Err(ClusteringError::InvalidInput(format!(
171            "Vectors must have the same length: got {} and {}",
172            a.len(),
173            b.len()
174        )));
175    }
176    if a.is_empty() {
177        return Ok(F::zero());
178    }
179    let default_config = SimdClusterConfig::default();
180    let cfg = config.unwrap_or(&default_config);
181
182    if should_use_simd(a.len(), cfg) {
183        let diff = F::simd_sub(&a, &b);
184        let abs_diff = F::simd_abs(&diff.view());
185        Ok(F::simd_sum(&abs_diff.view()))
186    } else {
187        let mut sum = F::zero();
188        for i in 0..a.len() {
189            sum = sum + (a[i] - b[i]).abs();
190        }
191        Ok(sum)
192    }
193}
194
195/// Compute distance between two vectors using a specified metric.
196///
197/// Dispatches to the appropriate SIMD-accelerated distance function.
198///
199/// # Errors
200///
201/// Returns `ClusteringError::InvalidInput` if vectors have different lengths.
202pub fn simd_distance<F>(
203    a: ArrayView1<F>,
204    b: ArrayView1<F>,
205    metric: SimdDistanceMetric,
206    config: Option<&SimdClusterConfig>,
207) -> Result<F>
208where
209    F: Float + FromPrimitive + Debug + SimdUnifiedOps,
210{
211    match metric {
212        SimdDistanceMetric::Euclidean => simd_euclidean_distance(a, b, config),
213        SimdDistanceMetric::SquaredEuclidean => simd_squared_euclidean_distance(a, b, config),
214        SimdDistanceMetric::Manhattan => simd_manhattan_distance(a, b, config),
215    }
216}
217
218// ---------------------------------------------------------------------------
219// Pairwise Distance Matrix
220// ---------------------------------------------------------------------------
221
222/// Compute the full pairwise distance matrix for a dataset using SIMD.
223///
224/// Returns an n x n symmetric distance matrix. For large datasets,
225/// consider using [`simd_pairwise_condensed_distances`] instead to save memory.
226///
227/// # Arguments
228///
229/// * `data` - Input data matrix (n_samples x n_features)
230/// * `metric` - Distance metric to use
231/// * `config` - Optional SIMD configuration
232///
233/// # Returns
234///
235/// Symmetric distance matrix of shape (n_samples x n_samples).
236///
237/// # Errors
238///
239/// Returns error if data is empty.
240pub fn simd_pairwise_distance_matrix<F>(
241    data: ArrayView2<F>,
242    metric: SimdDistanceMetric,
243    config: Option<&SimdClusterConfig>,
244) -> Result<Array2<F>>
245where
246    F: Float + FromPrimitive + Debug + Send + Sync + SimdUnifiedOps,
247{
248    let n_samples = data.shape()[0];
249    if n_samples == 0 {
250        return Err(ClusteringError::InvalidInput(
251            "Data must have at least one sample".to_string(),
252        ));
253    }
254
255    let default_config = SimdClusterConfig::default();
256    let cfg = config.unwrap_or(&default_config);
257    let mut distances = Array2::zeros((n_samples, n_samples));
258
259    let use_parallel =
260        cfg.enable_parallel && is_parallel_enabled() && n_samples > cfg.parallel_chunk_size;
261
262    if use_parallel {
263        simd_pairwise_matrix_parallel(data, metric, cfg, &mut distances);
264    } else {
265        simd_pairwise_matrix_sequential(data, metric, cfg, &mut distances)?;
266    }
267
268    Ok(distances)
269}
270
271/// Sequential pairwise distance matrix computation with SIMD.
272fn simd_pairwise_matrix_sequential<F>(
273    data: ArrayView2<F>,
274    metric: SimdDistanceMetric,
275    config: &SimdClusterConfig,
276    distances: &mut Array2<F>,
277) -> Result<()>
278where
279    F: Float + FromPrimitive + Debug + SimdUnifiedOps,
280{
281    let n_samples = data.shape()[0];
282    let block_size = config.block_size;
283
284    for block_i in (0..n_samples).step_by(block_size) {
285        let end_i = (block_i + block_size).min(n_samples);
286        for block_j in (block_i..n_samples).step_by(block_size) {
287            let end_j = (block_j + block_size).min(n_samples);
288
289            for i in block_i..end_i {
290                let start_j = if block_i == block_j { i + 1 } else { block_j };
291                for j in start_j..end_j {
292                    let dist = simd_distance(data.row(i), data.row(j), metric, Some(config))?;
293                    distances[[i, j]] = dist;
294                    distances[[j, i]] = dist;
295                }
296            }
297        }
298    }
299    Ok(())
300}
301
302/// Parallel pairwise distance matrix computation with SIMD.
303fn simd_pairwise_matrix_parallel<F>(
304    data: ArrayView2<F>,
305    metric: SimdDistanceMetric,
306    config: &SimdClusterConfig,
307    distances: &mut Array2<F>,
308) where
309    F: Float + FromPrimitive + Debug + Send + Sync + SimdUnifiedOps,
310{
311    let n_samples = data.shape()[0];
312
313    // Compute upper triangle in parallel by row
314    let row_results: Vec<Vec<(usize, F)>> = (0..n_samples)
315        .into_par_iter()
316        .map(|i| {
317            let mut row_dists = Vec::with_capacity(n_samples - i - 1);
318            for j in (i + 1)..n_samples {
319                let dist = simd_distance(data.row(i), data.row(j), metric, Some(config))
320                    .unwrap_or_else(|_| F::zero());
321                row_dists.push((j, dist));
322            }
323            row_dists
324        })
325        .collect();
326
327    // Fill symmetric matrix from parallel results
328    for (i, row_dists) in row_results.into_iter().enumerate() {
329        for (j, dist) in row_dists {
330            distances[[i, j]] = dist;
331            distances[[j, i]] = dist;
332        }
333    }
334}
335
336/// Compute condensed pairwise distances (upper triangular, no diagonal).
337///
338/// This is the memory-efficient form: stores n*(n-1)/2 distances instead of n*n.
339/// The index mapping is: `condensed_idx = n * i + j - (i + 1) * (i + 2) / 2`
340/// for row i < j.
341///
342/// # Arguments
343///
344/// * `data` - Input data matrix (n_samples x n_features)
345/// * `metric` - Distance metric
346/// * `config` - Optional SIMD configuration
347///
348/// # Returns
349///
350/// Condensed distance vector of length n*(n-1)/2.
351///
352/// # Errors
353///
354/// Returns error if data is empty.
355pub fn simd_pairwise_condensed_distances<F>(
356    data: ArrayView2<F>,
357    metric: SimdDistanceMetric,
358    config: Option<&SimdClusterConfig>,
359) -> Result<Array1<F>>
360where
361    F: Float + FromPrimitive + Debug + Send + Sync + SimdUnifiedOps,
362{
363    let n_samples = data.shape()[0];
364    if n_samples == 0 {
365        return Err(ClusteringError::InvalidInput(
366            "Data must have at least one sample".to_string(),
367        ));
368    }
369
370    let n_distances = n_samples * (n_samples - 1) / 2;
371    let mut distances = Array1::zeros(n_distances);
372    let default_config = SimdClusterConfig::default();
373    let cfg = config.unwrap_or(&default_config);
374
375    let mut idx = 0;
376    for i in 0..n_samples {
377        for j in (i + 1)..n_samples {
378            let dist = simd_distance(data.row(i), data.row(j), metric, Some(cfg))?;
379            distances[idx] = dist;
380            idx += 1;
381        }
382    }
383
384    Ok(distances)
385}