koan_core/index/
features.rs1use std::path::Path;
9
10use thiserror::Error;
11
12pub 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
21pub 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
34pub 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
43pub 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
56pub 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
67pub 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 assert!(bytes_to_embedding(&[0u8; 10]).is_none());
104 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 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}