klib/ml/base/

helpers.rs

//! Module for generic training and inference helpers.

use std::{
    collections::hash_map::DefaultHasher,
    fs::File,
    hash::{Hash, Hasher},
    io::{BufReader, Cursor, Write},
    path::{Path, PathBuf},
};

use anyhow::Context;
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

use crate::core::{
    base::Res,
    helpers::{inv_mel, mel},
    note::{HasNoteId, Note, ALL_PITCH_NOTES_WITH_FREQUENCY},
    pitch::HasFrequency,
};

use super::{KordItem, FREQUENCY_SPACE_SIZE, MEL_SPACE_SIZE, NOTE_SIGNATURE_SIZE, PITCH_CLASS_COUNT};
#[cfg(feature = "ml_loader_frequency_pooled")]
use super::{FREQUENCY_POOL_FACTOR, FREQUENCY_SPACE_POOLED_SIZE};

// Operations for working with kord samples.

/// Load the kord sample from the binary file into a new [`KordItem`].
pub fn load_kord_item(path: impl AsRef<Path>) -> Res<KordItem> {
    let path = path.as_ref();
    let file = std::fs::File::open(path).with_context(|| format!("File: `{path:?}`."))?;
    let mut reader = BufReader::new(file);

    // Read 8192 f32s in big endian from the file.
    let mut frequency_space = [0f32; 8192];

    for k in 0..FREQUENCY_SPACE_SIZE {
        frequency_space[k] = reader.read_f32::<BigEndian>()?;
    }

    let label = reader.read_u128::<BigEndian>()?;

    Ok(KordItem {
        path: path.to_owned(),
        frequency_space,
        label,
    })
}

/// Save the kord sample into a binary file.
pub fn save_kord_item(destination: impl AsRef<Path>, prefix: &str, note_names: &str, item: &KordItem) -> Res<PathBuf> {
    let mut output_data: Vec<u8> = Vec::with_capacity(FREQUENCY_SPACE_SIZE);
    let mut cursor = Cursor::new(&mut output_data);

    // Write frequency space.
    for value in item.frequency_space {
        cursor.write_f32::<BigEndian>(value)?;
    }

    // Write result.
    cursor.write_u128::<BigEndian>(item.label)?;

    // Get the hash.
    let mut hasher = DefaultHasher::new();
    output_data.hash(&mut hasher);
    let hash = hasher.finish();

    // Write the file.
    let path = destination.as_ref().join(format!("{prefix}{note_names}_{hash}.bin"));
    let mut f = File::create(&path)?;
    f.write_all(&output_data)?;

    Ok(path)
}

// Operations for working with mels.

/// Convert the [`FREQUENCY_SPACE_SIZE`] f32s in frequency space into [`MEL_SPACE_SIZE`] mel filter bands.
pub fn mel_filter_banks_from(spectrum: &[f32]) -> [f32; MEL_SPACE_SIZE] {
    let num_frequencies = spectrum.len();
    let num_mels = MEL_SPACE_SIZE;

    let f_min = 0f32;
    let f_max = FREQUENCY_SPACE_SIZE as f32;

    let mel_points = linspace(mel(f_min), mel(f_max), num_mels + 2);
    let f_points = mel_points.iter().map(|m| inv_mel(*m)).collect::<Vec<_>>();

    let mut filter_banks = [0f32; MEL_SPACE_SIZE];

    for i in 0..num_mels {
        let f_m_minus = f_points[i];
        let f_m = f_points[i + 1];
        let f_m_plus = f_points[i + 2];

        let k_minus = (num_frequencies as f32 * f_m_minus / 8192f32).floor() as usize;
        let k = (num_frequencies as f32 * f_m / 8192f32).floor() as usize;
        let k_plus = (num_frequencies as f32 * f_m_plus / 8192f32).floor() as usize;

        for j in k_minus..k {
            filter_banks[i] += spectrum[j] * (j - k_minus) as f32 / (k - k_minus) as f32;
        }

        for j in k..k_plus {
            filter_banks[i] += spectrum[j] * (k_plus - j) as f32 / (k_plus - k) as f32;
        }
    }

    filter_banks
}

/// Run a note-binned "harmonic convolution" over the frequency space data.
pub fn note_binned_convolution(spectrum: &[f32]) -> [f32; NOTE_SIGNATURE_SIZE] {
    let mut convolution = [0f32; NOTE_SIGNATURE_SIZE];

    for (note, _) in ALL_PITCH_NOTES_WITH_FREQUENCY.iter().skip(7).take(90) {
        let id_index = note.id_index();

        let (low, high) = note.tight_frequency_range();
        let low = low.round() as usize;
        let high = high.round() as usize;

        if high >= FREQUENCY_SPACE_SIZE {
            continue;
        }

        let mut sum = 0f32;
        for k in low..high {
            sum += spectrum[k];
        }

        convolution[id_index as usize] = sum;
    }

    convolution
}

/// Run a "harmonic convolution" over the frequency space data.
pub fn harmonic_convolution(spectrum: &[f32]) -> [f32; FREQUENCY_SPACE_SIZE] {
    let mut harmonic_convolution = [0f32; FREQUENCY_SPACE_SIZE];

    let (peak, _) = spectrum.iter().enumerate().fold((0usize, 0f32), |(k, max), (j, x)| if *x > max { (j, *x) } else { (k, max) });

    for center in (peak / 2)..4000 {
        let mut sum = spectrum[center];

        for k in 2..16 {
            let index = center * k;
            if index < FREQUENCY_SPACE_SIZE {
                sum += spectrum[index];
            }
        }

        for k in 2..16 {
            let index = center / k;
            if index < FREQUENCY_SPACE_SIZE {
                sum -= spectrum[index];
            }
        }

        harmonic_convolution[center] = sum.clamp(0.0, f32::MAX);
    }

    harmonic_convolution
}

/// Downsamples the full frequency space by averaging contiguous windows.
#[cfg(feature = "ml_loader_frequency_pooled")]
pub fn average_pool_frequency_space(spectrum: &[f32; FREQUENCY_SPACE_SIZE]) -> [f32; FREQUENCY_SPACE_POOLED_SIZE] {
    let mut pooled = [0f32; FREQUENCY_SPACE_POOLED_SIZE];

    for (index, chunk) in spectrum.chunks_exact(FREQUENCY_POOL_FACTOR).enumerate() {
        let sum: f32 = chunk.iter().sum();
        pooled[index] = sum / FREQUENCY_POOL_FACTOR as f32;
    }

    pooled
}

/// Create a linearly spaced vector.
pub fn linspace(start: f32, end: f32, num_points: usize) -> Vec<f32> {
    let step = (end - start) / (num_points - 1) as f32;
    (0..num_points).map(|i| start + i as f32 * step).collect()
}

/// Gets the "deterministic guess" for a given kord item.
#[cfg(feature = "analyze_base")]
pub fn get_deterministic_guess(kord_item: &KordItem) -> u128 {
    use crate::analyze::base::get_notes_from_smoothed_frequency_space;

    let smoothed_frequency_space = kord_item.frequency_space.into_iter().enumerate().map(|(k, v)| (k as f32, v)).collect::<Vec<_>>();

    let notes = get_notes_from_smoothed_frequency_space(&smoothed_frequency_space);

    Note::id_mask(&notes)
}

/// Produces a 128 element array of 0s and 1s from a u128.
pub fn u128_to_binary(num: u128) -> [f32; 128] {
    let mut binary = [0f32; 128];
    for i in 0..128 {
        binary[128 - 1 - i] = (num >> i & 1) as f32;
    }

    binary
}

/// Produces a u128 from a 128 element array of 0s and 1s.
pub fn binary_to_u128(binary: &[f32]) -> u128 {
    let mut num = 0u128;
    for i in 0..128 {
        num += (binary[i] as u128) << (128 - 1 - i);
    }

    num
}

/// Folds the 128-bit binary signature of the the notes into a 12-bit signature (which represent one octave)
#[allow(dead_code)]
pub fn fold_binary(binary: &[f32; NOTE_SIGNATURE_SIZE]) -> [f32; PITCH_CLASS_COUNT] {
    let mut folded = [0f32; PITCH_CLASS_COUNT];

    // binary is MSB-first from u128_to_binary, so we need to map array indices back to bit positions
    for array_idx in 0..NOTE_SIGNATURE_SIZE {
        if binary[array_idx] == 1.0 {
            // Convert array index back to bit position: array_idx corresponds to bit (127 - array_idx)
            let bit_position = NOTE_SIGNATURE_SIZE - 1 - array_idx;
            let pitch_class = bit_position % PITCH_CLASS_COUNT;
            folded[pitch_class] = 1.0;
        }
    }

    folded
}

/// Applies sigmoid activation to convert logits to probabilities in `[0, 1]`.
#[cfg(any(feature = "ml_target_full", feature = "ml_target_folded"))]
pub fn logits_to_probabilities(logits: &[f32]) -> Vec<f32> {
    logits.iter().map(|&logit| 1.0 / (1.0 + (-logit).exp())).collect()
}

/// Applies sigmoid activation and bass softmax to convert logits to probabilities in `[0, 1]`.
#[cfg(feature = "ml_target_folded_bass")]
pub fn logits_to_probabilities(logits: &[f32]) -> Vec<f32> {
    let mut probabilities: Vec<f32> = logits.iter().map(|&logit| 1.0 / (1.0 + (-logit).exp())).collect();

    if logits.len() >= PITCH_CLASS_COUNT {
        let slice = &logits[..PITCH_CLASS_COUNT];
        let max_logit = slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let exp_values: Vec<f32> = slice.iter().map(|&value| (value - max_logit).exp()).collect();
        let sum: f32 = exp_values.iter().sum();

        if sum > 0.0 {
            for (offset, exp_value) in exp_values.iter().enumerate() {
                probabilities[offset] = exp_value / sum;
            }
        }
    }

    probabilities
}

/// Converts probabilities to binary predictions using per-class thresholds.
#[cfg(any(feature = "ml_target_full", feature = "ml_target_folded"))]
pub fn logits_to_predictions(probabilities: &[f32], thresholds: &[f32]) -> Vec<f32> {
    probabilities
        .iter()
        .enumerate()
        .map(|(idx, probability)| {
            let threshold = thresholds.get(idx).copied().unwrap_or(0.5);
            if *probability > threshold {
                1.0
            } else {
                0.0
            }
        })
        .collect()
}

/// Converts probabilities to binary predictions using per-class thresholds.
#[cfg(feature = "ml_target_folded_bass")]
pub fn logits_to_predictions(probabilities: &[f32], thresholds: &[f32]) -> Vec<f32> {
    let mut predictions = vec![0.0; probabilities.len()];

    if probabilities.len() >= PITCH_CLASS_COUNT {
        let bass_slice = &probabilities[..PITCH_CLASS_COUNT];
        if let Some((best_idx, _)) = bass_slice.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) {
            predictions[best_idx] = 1.0;
        }

        for offset in 0..PITCH_CLASS_COUNT {
            let idx = PITCH_CLASS_COUNT + offset;
            if idx < probabilities.len() {
                let threshold = thresholds.get(idx).copied().unwrap_or(0.5);
                let probability = probabilities[idx];
                predictions[idx] = if probability > threshold { 1.0 } else { 0.0 };
            }
        }
    }

    predictions
}