Skip to main content

koan_core/index/
features.rs

1//! Acoustic feature extraction using bliss-audio.
2//!
3//! Extracts a feature vector per track (tempo, timbre, chroma,
4//! spectral features) for acoustic similarity search, with as many
5//! dimensions as bliss-audio's `NUMBER_FEATURES`.
6
7use std::path::Path;
8
9use thiserror::Error;
10
11/// Number of dimensions in the acoustic feature vector (matches bliss-audio).
12pub const EMBEDDING_DIMS: usize = bliss_audio::NUMBER_FEATURES;
13
14#[derive(Debug, Error)]
15pub enum AnalysisError {
16    #[error("bliss analysis failed: {0}")]
17    Bliss(String),
18}
19
20/// Analyze a track and return its acoustic feature vector.
21///
22/// Uses bliss-audio's Symphonia-based decoder to extract tempo, timbre,
23/// chroma, and spectral features. Takes ~0.4s per track on average.
24pub fn analyze_track(path: &Path) -> Result<Vec<f32>, AnalysisError> {
25    use bliss_audio::decoder::Decoder as _;
26    use bliss_audio::decoder::symphonia::SymphoniaDecoder;
27
28    let song =
29        SymphoniaDecoder::song_from_path(path).map_err(|e| AnalysisError::Bliss(e.to_string()))?;
30    Ok(song.analysis.as_vec())
31}
32
33/// Serialize a float vector to bytes for BLOB storage. Little-endian f32.
34pub fn embedding_to_bytes(embedding: &[f32]) -> Vec<u8> {
35    let mut bytes = Vec::with_capacity(embedding.len() * 4);
36    for &val in embedding {
37        bytes.extend_from_slice(&val.to_le_bytes());
38    }
39    bytes
40}
41
42/// Deserialize bytes from BLOB storage back to a float vector.
43pub fn bytes_to_embedding(bytes: &[u8]) -> Option<Vec<f32>> {
44    if !bytes.len().is_multiple_of(4) {
45        return None;
46    }
47    let count = bytes.len() / 4;
48    let mut embedding = Vec::with_capacity(count);
49    for chunk in bytes.as_chunks::<4>().0 {
50        embedding.push(f32::from_le_bytes(*chunk));
51    }
52    Some(embedding)
53}
54
55/// Compute euclidean distance between two embedding vectors.
56/// Vectors must be the same length.
57pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
58    debug_assert_eq!(a.len(), b.len());
59    a.iter()
60        .zip(b.iter())
61        .map(|(x, y)| (x - y) * (x - y))
62        .sum::<f32>()
63        .sqrt()
64}
65
66/// Compute the centroid (mean) of multiple embedding vectors.
67pub fn centroid(embeddings: &[Vec<f32>]) -> Vec<f32> {
68    if embeddings.is_empty() {
69        return vec![0.0; EMBEDDING_DIMS];
70    }
71    let dims = embeddings[0].len();
72    let mut result = vec![0.0f32; dims];
73    let count = embeddings.len() as f32;
74    for emb in embeddings {
75        for (i, &val) in emb.iter().enumerate() {
76            result[i] += val;
77        }
78    }
79    for val in &mut result {
80        *val /= count;
81    }
82    result
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn embedding_serialization_roundtrip() {
91        let embedding: Vec<f32> = (0..EMBEDDING_DIMS)
92            .map(|i| (i as f32) * 1.5 - 10.0)
93            .collect();
94        let bytes = embedding_to_bytes(&embedding);
95        let recovered = bytes_to_embedding(&bytes).unwrap();
96        assert_eq!(embedding, recovered);
97    }
98
99    #[test]
100    fn bytes_to_embedding_wrong_length() {
101        // Not a multiple of 4
102        assert!(bytes_to_embedding(&[0u8; 10]).is_none());
103        // Empty is valid (0 floats)
104        assert_eq!(bytes_to_embedding(&[]).unwrap().len(), 0);
105    }
106
107    #[test]
108    fn euclidean_distance_identical() {
109        let a = vec![1.0f32; EMBEDDING_DIMS];
110        assert!((euclidean_distance(&a, &a) - 0.0).abs() < f32::EPSILON);
111    }
112
113    #[test]
114    fn euclidean_distance_known() {
115        let mut a = vec![0.0f32; EMBEDDING_DIMS];
116        let mut b = vec![0.0f32; EMBEDDING_DIMS];
117        a[0] = 3.0;
118        b[0] = 0.0;
119        a[1] = 0.0;
120        b[1] = 4.0;
121        // sqrt(9 + 16) = 5.0
122        assert!((euclidean_distance(&a, &b) - 5.0).abs() < 1e-6);
123    }
124
125    #[test]
126    fn centroid_single() {
127        let emb = vec![42.0f32; EMBEDDING_DIMS];
128        let result = centroid(std::slice::from_ref(&emb));
129        assert_eq!(result, emb);
130    }
131
132    #[test]
133    fn centroid_multiple() {
134        let a = vec![2.0f32; EMBEDDING_DIMS];
135        let b = vec![4.0f32; EMBEDDING_DIMS];
136        let result = centroid(&[a, b]);
137        for val in result {
138            assert!((val - 3.0).abs() < f32::EPSILON);
139        }
140    }
141
142    #[test]
143    fn centroid_empty() {
144        let result = centroid(&[]);
145        assert_eq!(result.len(), EMBEDDING_DIMS);
146        for val in result {
147            assert!((val - 0.0).abs() < f32::EPSILON);
148        }
149    }
150}