use ndarray::{Array2, ArrayView2};
use rayon::prelude::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PerformanceHint {
CacheFriendly,
Vectorize,
Parallel,
GpuAccelerated,
LowLatency,
HighThroughput,
}
#[derive(Debug, Clone, Default)]
pub struct PerformanceHints {
hints: Vec<PerformanceHint>,
}
impl PerformanceHints {
pub fn new() -> Self {
Self::default()
}
pub fn with_hint(mut self, hint: PerformanceHint) -> Self {
self.hints.push(hint);
self
}
pub fn has_hint(&self, hint: PerformanceHint) -> bool {
self.hints.contains(&hint)
}
pub fn hints(&self) -> &[PerformanceHint] {
&self.hints
}
}
pub struct AudioKNN {
k_neighbors: usize,
}
impl AudioKNN {
pub fn new(k_neighbors: usize) -> Self {
Self { k_neighbors }
}
pub fn find_neighbors(
&self,
query: ArrayView2<f32>,
data: ArrayView2<f32>,
) -> Vec<Vec<(usize, f32)>> {
let n_queries = query.nrows();
let n_data = data.nrows();
(0..n_queries)
.into_par_iter()
.map(|q_idx| {
let query_point = query.row(q_idx);
let mut distances: Vec<(usize, f32)> = (0..n_data)
.map(|d_idx| {
let data_point = data.row(d_idx);
let dist = Self::euclidean_distance(query_point, data_point);
(d_idx, dist)
})
.collect();
distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
distances.truncate(self.k_neighbors);
distances
})
.collect()
}
#[inline(always)]
fn euclidean_distance(a: ndarray::ArrayView1<f32>, b: ndarray::ArrayView1<f32>) -> f32 {
a.iter()
.zip(b.iter())
.map(|(x, y)| {
let diff = x - y;
diff * diff
})
.sum::<f32>()
.sqrt()
}
pub fn find_similar_segments(
&self,
segment: ArrayView2<f32>,
audio: ArrayView2<f32>,
hop_size: usize,
) -> Vec<(usize, f32)> {
let segment_len = segment.ncols();
let audio_len = audio.ncols();
if segment_len > audio_len {
return Vec::new();
}
let num_windows = (audio_len - segment_len) / hop_size + 1;
let mut similarities: Vec<(usize, f32)> = (0..num_windows)
.into_par_iter()
.map(|i| {
let start = i * hop_size;
let end = start + segment_len;
let window = audio.slice(ndarray::s![.., start..end]);
let distance = Self::matrix_distance(&segment, &window);
(start, distance)
})
.collect();
similarities.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
similarities.truncate(self.k_neighbors);
similarities
}
fn matrix_distance(a: &ArrayView2<f32>, b: &ArrayView2<f32>) -> f32 {
a.iter()
.zip(b.iter())
.map(|(x, y)| {
let diff = x - y;
diff * diff
})
.sum::<f32>()
.sqrt()
}
}
pub struct BatchProcessor {
#[allow(dead_code)]
batch_size: usize,
num_threads: Option<usize>,
}
impl BatchProcessor {
pub fn new(batch_size: usize) -> Self {
Self {
batch_size,
num_threads: None,
}
}
pub fn with_threads(mut self, num_threads: usize) -> Self {
self.num_threads = Some(num_threads);
self
}
pub fn process<T, F, R>(&self, items: Vec<T>, f: F) -> Vec<R>
where
T: Send,
F: Fn(T) -> R + Send + Sync,
R: Send,
{
if let Some(threads) = self.num_threads {
rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.unwrap()
.install(|| items.into_par_iter().map(f).collect())
} else {
items.into_par_iter().map(f).collect()
}
}
pub fn process_chunks<F>(&self, data: &Array2<f32>, chunk_size: usize, f: F) -> Vec<Array2<f32>>
where
F: Fn(ArrayView2<f32>) -> Array2<f32> + Send + Sync,
{
let n_samples = data.ncols();
let n_chunks = n_samples.div_ceil(chunk_size);
(0..n_chunks)
.into_par_iter()
.map(|i| {
let start = i * chunk_size;
let end = (start + chunk_size).min(n_samples);
let chunk = data.slice(ndarray::s![.., start..end]);
f(chunk)
})
.collect()
}
}
pub struct SimdOps;
impl SimdOps {
#[inline]
pub fn multiply(a: &mut [f32], b: &[f32]) {
assert_eq!(a.len(), b.len());
a.iter_mut().zip(b.iter()).for_each(|(x, y)| *x *= y);
}
#[inline]
pub fn add(a: &mut [f32], b: &[f32]) {
assert_eq!(a.len(), b.len());
a.iter_mut().zip(b.iter()).for_each(|(x, y)| *x += y);
}
#[inline]
pub fn dot(a: &[f32], b: &[f32]) -> f32 {
assert_eq!(a.len(), b.len());
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}
#[inline]
pub fn rms(data: &[f32]) -> f32 {
let sum_squares: f32 = data.iter().map(|x| x * x).sum();
(sum_squares / data.len() as f32).sqrt()
}
pub fn find_peaks(data: &[f32], threshold: f32) -> Vec<usize> {
data.windows(3)
.enumerate()
.filter_map(|(i, window)| {
if window[1] > threshold && window[1] > window[0] && window[1] > window[2] {
Some(i + 1)
} else {
None
}
})
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_abs_diff_eq;
use ndarray::Array2;
#[test]
fn test_performance_hints() {
let hints = PerformanceHints::new()
.with_hint(PerformanceHint::Parallel)
.with_hint(PerformanceHint::Vectorize);
assert!(hints.has_hint(PerformanceHint::Parallel));
assert!(hints.has_hint(PerformanceHint::Vectorize));
assert!(!hints.has_hint(PerformanceHint::GpuAccelerated));
}
#[test]
fn test_audio_knn() {
let data = Array2::from_shape_vec(
(5, 3),
vec![
1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0, 4.0, 5.0, 6.0, 5.0, 6.0, 7.0,
],
)
.unwrap();
let query = Array2::from_shape_vec((1, 3), vec![2.5, 3.5, 4.5]).unwrap();
let knn = AudioKNN::new(3);
let neighbors = knn.find_neighbors(query.view(), data.view());
assert_eq!(neighbors.len(), 1);
assert_eq!(neighbors[0].len(), 3);
}
#[test]
fn test_simd_ops() {
let mut a = vec![1.0, 2.0, 3.0, 4.0];
let b = vec![2.0, 2.0, 2.0, 2.0];
SimdOps::multiply(&mut a, &b);
assert_eq!(a, vec![2.0, 4.0, 6.0, 8.0]);
let rms = SimdOps::rms(&[1.0, 2.0, 3.0, 4.0]);
assert_abs_diff_eq!(rms, 2.7386, epsilon = 0.001);
}
#[test]
fn test_batch_processor() {
let processor = BatchProcessor::new(10);
let items: Vec<i32> = (0..100).collect();
let results = processor.process(items, |x| x * 2);
assert_eq!(results.len(), 100);
assert_eq!(results[0], 0);
assert_eq!(results[99], 198);
}
}