koan_core/index/
features.rs1use std::path::Path;
8
9use thiserror::Error;
10
11pub 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
20pub 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
33pub 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
42pub 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
55pub 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
66pub 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 assert!(bytes_to_embedding(&[0u8; 10]).is_none());
103 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 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}