klib/ml/base/
data.rs

1//! Generic data structures and functions for training or inference.
2
3use burn::tensor::{backend::Backend, Tensor, TensorData};
4
5use super::{helpers::u128_to_binary, KordItem, INPUT_SPACE_SIZE, TARGET_SPACE_SIZE};
6
7#[cfg(feature = "ml_target_folded_bass")]
8use super::PITCH_CLASS_COUNT;
9
10#[cfg(any(feature = "ml_target_folded", feature = "ml_target_folded_bass"))]
11use super::helpers::fold_binary;
12
13#[cfg(feature = "ml_loader_note_binned_convolution")]
14use super::helpers::note_binned_convolution;
15
16#[cfg(feature = "ml_loader_mel")]
17use super::helpers::mel_filter_banks_from;
18
19#[cfg(feature = "ml_loader_frequency_pooled")]
20use super::helpers::average_pool_frequency_space;
21
22#[cfg(feature = "ml_loader_include_deterministic_guess")]
23use super::{helpers::get_deterministic_guess, DETERMINISTIC_GUESS_SIZE};
24
25/// Takes a loaded kord item and converts it to a sample tensor that is ready for classification.
26pub fn kord_item_to_sample_tensor<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
27    #[cfg(feature = "ml_loader_note_binned_convolution")]
28    {
29        sample_from_note_binned_convolution(device, item)
30    }
31
32    #[cfg(feature = "ml_loader_mel")]
33    {
34        sample_from_mel(device, item)
35    }
36
37    #[cfg(feature = "ml_loader_frequency")]
38    {
39        sample_from_frequency(device, item)
40    }
41
42    #[cfg(feature = "ml_loader_frequency_pooled")]
43    {
44        sample_from_frequency_pooled(device, item)
45    }
46}
47
48#[cfg(feature = "ml_loader_note_binned_convolution")]
49fn sample_from_note_binned_convolution<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
50    let mut convolution = note_binned_convolution(&item.frequency_space);
51    normalize(&mut convolution);
52    sample_tensor_from_parts(device, item, convolution.to_vec())
53}
54
55#[cfg(feature = "ml_loader_mel")]
56fn sample_from_mel<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
57    let mut mel_space = mel_filter_banks_from(&item.frequency_space);
58    normalize(&mut mel_space);
59    sample_tensor_from_parts(device, item, mel_space.to_vec())
60}
61
62#[cfg(feature = "ml_loader_frequency")]
63fn sample_from_frequency<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
64    let mut frequency_space = item.frequency_space;
65    normalize(&mut frequency_space);
66    sample_tensor_from_parts(device, item, frequency_space.to_vec())
67}
68
69#[cfg(feature = "ml_loader_frequency_pooled")]
70fn sample_from_frequency_pooled<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
71    let mut frequency_space = average_pool_frequency_space(&item.frequency_space);
72    normalize(&mut frequency_space);
73    sample_tensor_from_parts(device, item, frequency_space.to_vec())
74}
75
76#[allow(unused_variables)]
77fn sample_tensor_from_parts<B: Backend>(device: &B::Device, item: &KordItem, mut features: Vec<f32>) -> Tensor<B, 2> {
78    #[cfg(feature = "ml_loader_include_deterministic_guess")]
79    {
80        let mut combined = Vec::with_capacity(DETERMINISTIC_GUESS_SIZE + features.len());
81        combined.extend_from_slice(&deterministic_guess_array(item));
82        combined.extend_from_slice(&features);
83        features = combined;
84    }
85
86    to_zero_mean_unit_variance(features.as_mut_slice());
87
88    tensor_from_vec_with_expected_size(device, features, INPUT_SPACE_SIZE)
89}
90
91/// Takes a loaded kord item and converts it to a target tensor that is ready for classification.
92#[cfg(feature = "ml_target_full")]
93pub fn kord_item_to_target_tensor<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
94    let binary_full = u128_to_binary(item.label);
95    tensor_from_vec_with_expected_size(device, binary_full.to_vec(), TARGET_SPACE_SIZE)
96}
97
98/// Takes a loaded kord item and converts it to a folded target tensor that is ready for classification.
99#[cfg(feature = "ml_target_folded")]
100pub fn kord_item_to_target_tensor<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
101    let binary_full = u128_to_binary(item.label);
102    let folded = fold_binary(&binary_full);
103    tensor_from_vec_with_expected_size(device, folded.to_vec(), TARGET_SPACE_SIZE)
104}
105
106/// Takes a loaded kord item and converts it to a folded+bass target tensor that is ready for classification.
107#[cfg(feature = "ml_target_folded_bass")]
108pub fn kord_item_to_target_tensor<B: Backend>(device: &B::Device, item: &KordItem) -> Tensor<B, 2> {
109    let binary_full = u128_to_binary(item.label);
110    let folded = fold_binary(&binary_full);
111    let bass = lowest_pitch_class_mask(item.label);
112
113    let mut components = Vec::with_capacity(TARGET_SPACE_SIZE);
114    components.extend_from_slice(&bass);
115    components.extend_from_slice(&folded);
116
117    tensor_from_vec_with_expected_size(device, components, TARGET_SPACE_SIZE)
118}
119
120/// Modifies a slice in place to convert values to zero mean and unit variance.
121pub fn to_zero_mean_unit_variance(slice: &mut [f32]) {
122    let mean = slice.iter().sum::<f32>() / slice.len() as f32;
123    let variance = slice.iter().map(|x| (x - mean).powf(2.0)).sum::<f32>() / slice.len() as f32;
124    let std = variance.sqrt();
125
126    if std == 0.0 {
127        slice.fill(0.0);
128    } else {
129        slice.iter_mut().for_each(|x| *x = (*x - mean) / std);
130    }
131}
132
133/// Normalizes a slice in place.
134pub fn normalize(slice: &mut [f32]) {
135    let max = slice.iter().fold(0f32, |acc, &x| acc.max(x));
136
137    if max == 0.0 {
138        return;
139    }
140
141    slice.iter_mut().for_each(|x| *x /= max);
142}
143
144fn tensor_from_vec_with_expected_size<B: Backend>(device: &B::Device, data: Vec<f32>, expected: usize) -> Tensor<B, 2> {
145    debug_assert_eq!(data.len(), expected, "Tensor length mismatch: expected {expected}, received {}", data.len());
146
147    let len = data.len();
148    let tensor_data = TensorData::from(data.as_slice()).convert::<B::FloatElem>();
149    let tensor = Tensor::<B, 1>::from_data(tensor_data, device);
150
151    tensor.reshape([1, len])
152}
153
154#[cfg(feature = "ml_loader_include_deterministic_guess")]
155fn deterministic_guess_array(item: &KordItem) -> [f32; DETERMINISTIC_GUESS_SIZE] {
156    let guess_binary = get_deterministic_guess(item);
157    u128_to_binary(guess_binary)
158}
159
160#[cfg(feature = "ml_target_folded_bass")]
161fn lowest_pitch_class_mask(label: u128) -> [f32; PITCH_CLASS_COUNT] {
162    let mut mask = [0.0; PITCH_CLASS_COUNT];
163
164    if label != 0 {
165        let lowest = label.trailing_zeros() as usize;
166        let pitch_class = lowest % PITCH_CLASS_COUNT;
167        mask[pitch_class] = 1.0;
168    }
169
170    mask
171}
172
173#[cfg(all(test, feature = "ml_target_folded_bass"))]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn lowest_pitch_class_mask_marks_expected_bin() {
179        let label = 1u128 << 17; // arbitrary note index 17 -> pitch class 5
180        let mask = lowest_pitch_class_mask(label);
181
182        assert_eq!(mask.iter().sum::<f32>(), 1.0);
183        assert_eq!(mask[5], 1.0);
184    }
185
186    #[test]
187    fn lowest_pitch_class_mask_handles_empty_label() {
188        let mask = lowest_pitch_class_mask(0);
189        assert!(mask.iter().all(|v| *v == 0.0));
190    }
191}
klib/ml/base/data.rs

klib/ml/base/
data.rs
