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