diff --git a/Cargo.lock b/Cargo.lock index d648610..3ea28e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -54,7 +54,7 @@ dependencies = [ [[package]] name = "audio-pattern-detector" -version = "0.4.1" +version = "0.4.2" dependencies = [ "audio-pattern-detector-core", "clap", @@ -64,7 +64,7 @@ dependencies = [ [[package]] name = "audio-pattern-detector-core" -version = "0.4.1" +version = "0.4.2" dependencies = [ "base64", "pyo3", diff --git a/Cargo.toml b/Cargo.toml index c9dfd60..b8bd3db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ resolver = "2" members = ["crates/core", "crates/cli"] [workspace.package] -version = "0.4.1" +version = "0.4.2" edition = "2021" authors = ["Andrew Chen "] diff --git a/README.md b/README.md index 68fd38c..4b99e09 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ cargo install --git https://github.com/andrewtheguy/audio_pattern_detector.git - `ghcr.io/andrewtheguy/audio_pattern_detector:` is published with each release (`v`, plus `latest` for releases from `main`). It contains only the binary at `/usr/local/bin/audio-pattern-detector` (no ffmpeg), so the usual use is copying it into another image: ```dockerfile -COPY --from=ghcr.io/andrewtheguy/audio_pattern_detector:v0.4.1 /usr/local/bin/audio-pattern-detector /usr/local/bin/audio-pattern-detector +COPY --from=ghcr.io/andrewtheguy/audio_pattern_detector:v0.4.2 /usr/local/bin/audio-pattern-detector /usr/local/bin/audio-pattern-detector ``` The release workflow builds the image for amd64 and arm64 from the binaries it has already compiled (`runtime-prebuilt` target). To build it locally from source for the host architecture: `docker build --target runtime -t audio-pattern-detector .` diff --git a/crates/core/src/detector.rs b/crates/core/src/detector.rs index 626f84f..37c579c 100644 --- a/crates/core/src/detector.rs +++ b/crates/core/src/detector.rs @@ -10,7 +10,7 @@ use crate::audio_clip::{ }; use crate::dsp::{ fft_correlate_full, find_peaks_1d, integrated_loudness, loudness_normalize, pearson_correlation_1d, - resample_preserve_maxima_1d, FindPeaksOptions, + resample_preserve_maxima_1d, CorrelationTemplate, CorrelationWorkspace, FindPeaksOptions, }; use crate::error::{Error, Result}; use crate::stream::AudioStream; @@ -173,23 +173,66 @@ fn mean(values: &[f32]) -> f32 { (values.iter().map(|&v| v as f64).sum::() / values.len() as f64) as f32 } -/// Loudness-normalize audio to -16 dB LUFS. Silence comes out as NaN. -fn normalize_loudness(audio: &[f32], sample_rate: u32) -> Vec { +/// Loudness-normalize audio in place to -16 dB LUFS. +fn normalize_loudness(audio: &mut [f32], sample_rate: u32) { let seconds = audio.len() as f64 / sample_rate as f64; let block_size = if seconds < 0.5 { seconds } else { 0.4 }; let loudness = integrated_loudness(audio, sample_rate, block_size); - loudness_normalize(audio, loudness, TARGET_LUFS) + loudness_normalize(audio, loudness, TARGET_LUFS); +} + +/// Replace every value with its absolute value and return the largest. +fn absolute_in_place(values: &mut [f32]) -> f32 { + let mut absolute_max = f32::NEG_INFINITY; + for v in values.iter_mut() { + *v = v.abs(); + absolute_max = absolute_max.max(*v); + } + absolute_max } /// Absolute self-correlation normalized to a peak of 1, plus the original peak. fn clip_correlation(clip: &[f32]) -> (Vec, f32) { let mut correlation = fft_correlate_full(clip, clip); - correlation.iter_mut().for_each(|v| *v = v.abs()); - let absolute_max = correlation.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let absolute_max = absolute_in_place(&mut correlation); correlation.iter_mut().for_each(|v| *v /= absolute_max); (correlation, absolute_max) } +/// Clips whose sections are identical (same lookback), so one +/// loudness-normalized section and one forward FFT serve all of them. +struct SectionGroup { + /// Seconds of the previous chunk prepended to every chunk after the + /// first: ceil(clip seconds), the same for every clip in the group. + lookback_seconds: u32, + /// Indices into `AudioPatternDetector::clips`. + clip_indices: Vec, +} + +/// Group clips by sliding window, in order of first appearance. +fn section_groups(clips: &[ClipData]) -> Vec { + let mut groups: Vec = Vec::new(); + for (clip_index, clip_data) in clips.iter().enumerate() { + match groups.iter_mut().find(|g| g.lookback_seconds == clip_data.sliding_window) { + Some(group) => group.clip_indices.push(clip_index), + None => groups.push(SectionGroup { + lookback_seconds: clip_data.sliding_window, + clip_indices: vec![clip_index], + }), + } + } + groups +} + +/// Buffers reused across the chunks of one [`AudioPatternDetector::find_clip_in_audio`] run. +struct RunBuffers { + workspace: CorrelationWorkspace, + /// One cached correlation template per clip, in clip order. + templates: Vec, + audio_section: Vec, + correlation: Vec, +} + /// Slice `width` samples centered on `middle_index`, zero-padding past either end. pub fn slicing_with_zero_padding(array: &[f32], width: usize, middle_index: usize) -> Vec { let begin = middle_index as isize - (width / 2) as isize; @@ -211,6 +254,8 @@ pub struct AudioPatternDetector { target_sample_rate: u32, seconds_per_chunk: u32, min_chunk_size: u32, + /// Clips sharing a section, grouped by sliding window. + section_groups: Vec, } impl AudioPatternDetector { @@ -273,10 +318,11 @@ impl AudioPatternDetector { debug_mode = false; } - let clips = audio_clips + let clips: Vec = audio_clips .into_iter() .map(|audio_clip| Self::prepare_clip(audio_clip, sr, debug_mode)) .collect(); + let section_groups = section_groups(&clips); Ok(Self { clips, @@ -286,6 +332,7 @@ impl AudioPatternDetector { target_sample_rate: sr, seconds_per_chunk, min_chunk_size, + section_groups, }) } @@ -300,7 +347,8 @@ impl AudioPatternDetector { ); } - let clip = normalize_loudness(&audio_clip.audio, sr); + let mut clip = audio_clip.audio; + normalize_loudness(&mut clip, sr); let (correlation_clip, correlation_clip_absolute_max) = clip_correlation(&clip); if debug_mode { @@ -416,6 +464,12 @@ impl AudioPatternDetector { // previous chunk. The next chunk sees the end of that audio again // through its lookback and would report the same matches twice. let mut previous_peaks: Vec> = vec![Vec::new(); self.clips.len()]; + let mut buffers = RunBuffers { + workspace: CorrelationWorkspace::new(), + templates: self.clips.iter().map(|c| CorrelationTemplate::new(&c.clip)).collect(), + audio_section: Vec::new(), + correlation: Vec::new(), + }; loop { let chunk = audio_stream.source.read_samples(chunk_samples)?; @@ -424,33 +478,39 @@ impl AudioPatternDetector { } total_time += chunk.len() as f64 / sr as f64; - // All matches from all clips for this chunk. - let mut chunk_matches: Vec<(f64, &str)> = Vec::new(); - - for (clip_data, seen_peaks) in self.clips.iter().zip(previous_peaks.iter_mut()) { - let peaks = self.process_chunk(&chunk, clip_data, previous_chunk.as_deref(), index)?; - let peak_times: Vec = peaks - .iter() - .filter(|(position, _)| !seen_peaks.contains(position)) - .map(|&(_, timestamp)| timestamp) - .collect(); - *seen_peaks = peaks.into_iter().map(|(position, _)| position).collect(); - - if on_pattern_detected.is_some() { - chunk_matches.extend(peak_times.iter().map(|&t| (t, clip_data.name.as_str()))); - } - if let Some(all) = all_peak_times.as_mut() { - all.get_mut(&clip_data.name) - .expect("every clip has an entry") - .extend(peak_times); + // All matches from all clips for this chunk, with the clip index. + let mut chunk_matches: Vec<(f64, usize)> = Vec::new(); + + for group in &self.section_groups { + let lookback_samples = self.load_section(&chunk, previous_chunk.as_deref(), group, &mut buffers); + for &clip_index in &group.clip_indices { + let clip_data = &self.clips[clip_index]; + let seen_peaks = &mut previous_peaks[clip_index]; + let peaks = self.process_section(clip_index, lookback_samples, index, &mut buffers)?; + let peak_times: Vec = peaks + .iter() + .filter(|(position, _)| !seen_peaks.contains(position)) + .map(|&(_, timestamp)| timestamp) + .collect(); + *seen_peaks = peaks.into_iter().map(|(position, _)| position).collect(); + + if on_pattern_detected.is_some() { + chunk_matches.extend(peak_times.iter().map(|&t| (t, clip_index))); + } + if let Some(all) = all_peak_times.as_mut() { + all.get_mut(&clip_data.name) + .expect("every clip has an entry") + .extend(peak_times); + } } } - // Call the callback in timestamp order (monotonic output). + // Call the callback in timestamp order (monotonic output), clip + // order within the same timestamp. if let Some(callback) = on_pattern_detected.as_mut() { - chunk_matches.sort_by(|a, b| a.0.total_cmp(&b.0)); - for (timestamp, clip_name) in chunk_matches { - callback(clip_name, timestamp); + chunk_matches.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1))); + for (timestamp, clip_index) in chunk_matches { + callback(&self.clips[clip_index].name, timestamp); } } @@ -461,47 +521,64 @@ impl AudioPatternDetector { Ok((all_peak_times, total_time)) } - /// Detect one clip in one chunk; returns, per match, the peak's sample - /// position from the beginning of the stream and the timestamp of the - /// clip start in seconds from the beginning of the stream. - fn process_chunk( + /// Build the audio section of a chunk for one group of clips into + /// `buffers`, loudness-normalize it and transform it for correlation. + /// Returns the number of lookback samples prepended. + /// + /// The last `lookback_seconds` (= ceil(clip seconds)) of `previous_chunk` + /// are prepended so a pattern that crosses the boundary is fully + /// contained in the section. This is applied uniformly to every + /// non-first chunk — including the final short chunk, whose own length + /// is not a reliable lookback. The section depends only on the group's + /// lookback, so a clip's detections never depend on other clips. + fn load_section( &self, chunk: &[f32], - clip_data: &ClipData, previous_chunk: Option<&[f32]>, + group: &SectionGroup, + buffers: &mut RunBuffers, + ) -> usize { + let sr = self.target_sample_rate; + let lookback_samples = + previous_chunk.map_or(0, |previous| (group.lookback_seconds as usize * sr as usize).min(previous.len())); + + let RunBuffers { workspace, audio_section, .. } = buffers; + audio_section.clear(); + if let Some(previous) = previous_chunk { + audio_section.extend_from_slice(&previous[previous.len() - lookback_samples..]); + } + audio_section.extend_from_slice(chunk); + + normalize_loudness(audio_section, sr); + + workspace.load_signal(audio_section); + lookback_samples + } + + /// Detect one clip in the loaded section; returns, per match, the peak's + /// sample position from the beginning of the stream and the timestamp of + /// the clip start in seconds from the beginning of the stream. + fn process_section( + &self, + clip_index: usize, + lookback_samples: usize, index: usize, + buffers: &mut RunBuffers, ) -> Result> { let sr = self.target_sample_rate; + let clip_data = &self.clips[clip_index]; let clip_seconds = clip_data.clip.len() as f64 / sr as f64; + let lookback_seconds = lookback_samples as f64 / sr as f64; - // Prepend the last sliding_window seconds (= ceil(clip_seconds)) of - // previous_chunk so a pattern that crosses the boundary is fully - // contained in audio_section. Applied uniformly to every non-first - // chunk — including the final short chunk, whose own length is not a - // reliable lookback. - let lookback_samples = audio_section_lookback(clip_data, previous_chunk, sr); - let (subtract_seconds, audio_section) = match previous_chunk { - Some(previous) => { - let mut section = Vec::with_capacity(lookback_samples + chunk.len()); - section.extend_from_slice(&previous[previous.len() - lookback_samples..]); - section.extend_from_slice(chunk); - (clip_data.sliding_window as f64, section) - } - // First chunk: no previous audio available. - None => (0.0, chunk.to_vec()), - }; - - let mut audio_section = normalize_loudness(&audio_section, sr); - // NaN comes from loudness normalization of silence. - audio_section.iter_mut().filter(|v| v.is_nan()).for_each(|v| *v = 0.0); - - let peaks = self.correlation_method(clip_data, &audio_section, index)?; + let RunBuffers { workspace, templates, audio_section, correlation } = buffers; + workspace.correlate(&mut templates[clip_index], correlation); + let peaks = self.correlation_method(clip_data, audio_section, correlation, index)?; let section_start = index as i64 * self.seconds_per_chunk as i64 * sr as i64 - lookback_samples as i64; Ok(peaks .into_iter() .map(|peak| { - let peak_time = peak as f64 / sr as f64 - subtract_seconds; + let peak_time = peak as f64 / sr as f64 - lookback_seconds; let from_beginning = peak_time + (index as f64 * self.seconds_per_chunk as f64); // Move the timestamp to be before the clip. (section_start + peak as i64, (from_beginning - clip_seconds).max(0.0)) @@ -509,18 +586,23 @@ impl AudioPatternDetector { .collect()) } - /// Step 1 (cross-correlation + peak finding) followed by Step 2 - /// (verification) for one audio section. Returns accepted peak indices. - fn correlation_method(&self, clip_data: &ClipData, audio_section: &[f32], index: usize) -> Result> { + /// Step 1 (peak finding on the cross-correlation of the section with + /// the clip, given in `correlation`) followed by Step 2 (verification) + /// for one audio section. Returns accepted peak indices. + fn correlation_method( + &self, + clip_data: &ClipData, + audio_section: &[f32], + correlation: &mut [f32], + index: usize, + ) -> Result> { let clip_name = &clip_data.name; let clip_length = clip_data.clip.len(); let correlation_clip_length = clip_data.correlation_clip.len(); - // Cross-correlate and normalize by the larger of the two peaks so a - // much softer section cannot look like a full-strength match. - let mut correlation = fft_correlate_full(audio_section, &clip_data.clip); - correlation.iter_mut().for_each(|v| *v = v.abs()); - let absolute_max = correlation.iter().copied().fold(f32::NEG_INFINITY, f32::max); + // Normalize by the larger of the two peaks so a much softer section + // cannot look like a full-strength match. + let absolute_max = absolute_in_place(correlation); let max_choose = clip_data.correlation_clip_absolute_max.max(absolute_max); correlation.iter_mut().for_each(|v| *v /= max_choose); @@ -533,7 +615,7 @@ impl AudioPatternDetector { // No repetition within the duration of the clip. The height is kept // low so weak candidates are not missed; they are verified below. let peaks = find_peaks_1d( - &correlation, + correlation, &FindPeaksOptions { height: Some(self.height_min), distance: Some(clip_length), prominence: None }, ); @@ -557,7 +639,7 @@ impl AudioPatternDetector { let accepted = match &clip_data.tone { Some(tone) => self.verify_marker_tone(tone, audio_section, peak, clip_length, §ion_ts), - None => self.verify_correlation_envelope(clip_data, &correlation, peak, §ion_ts, &mut debug), + None => self.verify_correlation_envelope(clip_data, correlation, peak, §ion_ts, &mut debug), }; if accepted { peaks_final.push(peak); @@ -774,11 +856,6 @@ pub fn safe_path_component(name: &str) -> String { component } -/// Samples of the previous chunk prepended to a chunk's audio section. -fn audio_section_lookback(clip_data: &ClipData, previous_chunk: Option<&[f32]>, sr: u32) -> usize { - previous_chunk.map_or(0, |previous| (clip_data.sliding_window as usize * sr as usize).min(previous.len())) -} - /// Analyze the candidate window plus the clip-length windows on either side. /// Returns `(matched, left flank, right flank)` metrics. pub fn analyze_tone_candidate_context( @@ -885,6 +962,82 @@ mod tests { assert_eq!(downsample_window(&curve, (5, 10, 12)), (13..25).map(|i| i as f32).collect::>()); } + #[test] + fn test_absolute_in_place() { + let mut values = [-3.0_f32, 1.0, -0.5, 2.0, -0.0]; + assert_eq!(absolute_in_place(&mut values), 3.0); + assert_eq!(values, [3.0, 1.0, 0.5, 2.0, 0.0]); + + let mut negative = [-0.25_f32, -4.0, -1.0]; + assert_eq!(absolute_in_place(&mut negative), 4.0); + assert_eq!(negative, [0.25, 4.0, 1.0]); + + assert_eq!(absolute_in_place(&mut []), f32::NEG_INFINITY); + } + + #[test] + fn test_clip_correlation_is_normalized_with_original_peak() { + let clip = [0.5_f32, -1.0, 0.25, 0.75]; + let energy: f32 = clip.iter().map(|v| v * v).sum(); + let (correlation, absolute_max) = clip_correlation(&clip); + assert_eq!(correlation.len(), 2 * clip.len() - 1); + assert!((absolute_max - energy).abs() < 1e-5, "{absolute_max} != {energy}"); + assert_eq!(correlation[clip.len() - 1], 1.0); + assert!(correlation.iter().all(|&v| (0.0..=1.0).contains(&v)), "{correlation:?}"); + } + + fn sine_clip(name: &str, seconds: f64) -> AudioClip { + let sr = DEFAULT_TARGET_SAMPLE_RATE; + let samples = (sr as f64 * seconds) as usize; + let audio = (0..samples) + .map(|i| (2.0 * std::f64::consts::PI * 1000.0 * i as f64 / sr as f64).sin() as f32) + .collect(); + AudioClip::new(name, audio, sr) + } + + #[test] + fn test_section_groups_by_sliding_window_in_order_of_first_appearance() { + let detector = AudioPatternDetector::new( + vec![ + sine_clip("short", 0.23), // window 1 + sine_clip("long", 2.5), // window 3 + sine_clip("medium", 0.9), // window 1 + sine_clip("longer", 3.0), // window 3 (exactly 3 seconds) + sine_clip("one", 1.0), // window 1 (exactly 1 second) + sine_clip("two", 1.2), // window 2 + ], + DetectorOptions::default(), + ) + .unwrap(); + + let groups = &detector.section_groups; + assert_eq!(groups.len(), 3); + + assert_eq!(groups[0].lookback_seconds, 1); + assert_eq!(groups[0].clip_indices, vec![0, 2, 4]); + + assert_eq!(groups[1].lookback_seconds, 3); + assert_eq!(groups[1].clip_indices, vec![1, 3]); + + assert_eq!(groups[2].lookback_seconds, 2); + assert_eq!(groups[2].clip_indices, vec![5]); + + // Every clip is in exactly one group. + let mut all: Vec = groups.iter().flat_map(|g| g.clip_indices.iter().copied()).collect(); + all.sort_unstable(); + assert_eq!(all, (0..detector.clips.len()).collect::>()); + for group in groups { + for &clip_index in &group.clip_indices { + assert_eq!(detector.clips[clip_index].sliding_window, group.lookback_seconds); + } + } + } + + #[test] + fn test_section_groups_empty() { + assert!(section_groups(&[]).is_empty()); + } + #[test] fn test_detector_config_serializes_clips_as_ordered_map() { let config = DetectorConfig { diff --git a/crates/core/src/dsp/correlate.rs b/crates/core/src/dsp/correlate.rs index d5cb30e..800ac99 100644 --- a/crates/core/src/dsp/correlate.rs +++ b/crates/core/src/dsp/correlate.rs @@ -1,5 +1,6 @@ // ── FFT cross-correlation ──────────────────────────────────────────── +use realfft::num_complex::Complex; use realfft::{ComplexToReal, RealFftPlanner, RealToComplex}; use std::cell::RefCell; use std::collections::VecDeque; @@ -39,51 +40,167 @@ fn get_fft_plans(fft_size: usize) -> FftPlans { }) } -/// Cross-correlate two real 1-D signals using FFT, returning the full -/// correlation (length `signal.len() + template.len() - 1`). -/// -/// Matches `scipy.signal.correlate(signal, template, mode="full")`: output -/// index `k` is the lag where `template[template.len() - 1]` aligns with -/// `signal[k]`. Returns an empty vector if either input is empty. -pub fn fft_correlate_full(signal: &[f32], template: &[f32]) -> Vec { - if signal.is_empty() || template.is_empty() { - return Vec::new(); +/// Number of FFT sizes a [`CorrelationTemplate`] keeps spectra for. A run +/// sees at most three: the first chunk (no lookback), full chunks and the +/// final short chunk. +const TEMPLATE_SPECTRUM_CACHE_CAPACITY: usize = 4; + +/// A correlation template with its time-reversed spectrum cached per FFT +/// size, so repeated correlations against new signals only transform the +/// signal. +pub struct CorrelationTemplate { + reversed: Vec, + spectra: Vec<(usize, Vec>)>, +} + +impl CorrelationTemplate { + pub fn new(template: &[f32]) -> Self { + Self { + // Time-reversing the template turns convolution into correlation. + reversed: template.iter().rev().copied().collect(), + spectra: Vec::new(), + } } - let output_len = signal.len() + template.len() - 1; - let fft_size = output_len.next_power_of_two(); + pub fn len(&self) -> usize { + self.reversed.len() + } + + pub fn is_empty(&self) -> bool { + self.reversed.is_empty() + } + + /// Spectrum of the zero-padded, reversed template at `fft_size`. + fn spectrum(&mut self, fft_size: usize) -> &[Complex] { + if let Some(index) = self.spectra.iter().position(|(size, _)| *size == fft_size) { + return &self.spectra[index].1; + } + assert!(fft_size >= self.reversed.len(), "FFT size must cover the template"); + + let (r2c, _) = get_fft_plans(fft_size); + let mut padded = vec![0.0_f32; fft_size]; + padded[..self.reversed.len()].copy_from_slice(&self.reversed); + let mut spectrum = r2c.make_output_vec(); + let mut scratch = r2c.make_scratch_vec(); + r2c.process_with_scratch(&mut padded, &mut spectrum, &mut scratch) + .expect("forward FFT buffers are sized by the plan"); + + if self.spectra.len() == TEMPLATE_SPECTRUM_CACHE_CAPACITY { + self.spectra.remove(0); + } + self.spectra.push((fft_size, spectrum)); + &self.spectra.last().expect("spectrum was just pushed").1 + } +} + +/// Reusable buffers for FFT cross-correlation of one signal against any +/// number of templates: the signal is set by [`load_signal`] and transformed +/// once per FFT size, so each [`correlate`] call with a template needing the +/// same size costs a spectrum product and one inverse FFT. +/// +/// The FFT size of a correlation depends only on the signal and that +/// template, never on the other templates, so the result is bit-identical +/// to correlating the pair on its own. +/// +/// [`load_signal`]: CorrelationWorkspace::load_signal +/// [`correlate`]: CorrelationWorkspace::correlate +#[derive(Default)] +pub struct CorrelationWorkspace { + signal: Vec, + /// FFT size the buffers are allocated for (0 before the first transform). + fft_size: usize, + /// Whether `signal_spectrum` is the transform of `signal` at `fft_size`. + spectrum_loaded: bool, + plans: Option, + padded_signal: Vec, + signal_spectrum: Vec>, + product: Vec>, + forward_scratch: Vec>, + inverse_scratch: Vec>, + time_domain: Vec, +} + +impl CorrelationWorkspace { + pub fn new() -> Self { + Self::default() + } - let mut padded_signal = vec![0.0_f32; fft_size]; - let mut padded_template = vec![0.0_f32; fft_size]; - padded_signal[..signal.len()].copy_from_slice(signal); - // Time-reversing the template turns convolution into correlation. - for (dst, &val) in padded_template.iter_mut().zip(template.iter().rev()) { - *dst = val; + /// Set the signal for the following [`correlate`](Self::correlate) calls. + pub fn load_signal(&mut self, signal: &[f32]) { + self.signal.clear(); + self.signal.extend_from_slice(signal); + self.spectrum_loaded = false; } - let (r2c, c2r) = get_fft_plans(fft_size); - let mut signal_spectrum = r2c.make_output_vec(); - let mut template_spectrum = r2c.make_output_vec(); - let mut forward_scratch = r2c.make_scratch_vec(); + /// Transform the loaded signal at `fft_size` unless that is already done. + fn transform_signal(&mut self, fft_size: usize) { + if self.spectrum_loaded && self.fft_size == fft_size { + return; + } + if self.fft_size != fft_size || self.plans.is_none() { + let plans = get_fft_plans(fft_size); + let (r2c, c2r) = &plans; + self.signal_spectrum = r2c.make_output_vec(); + self.product = r2c.make_output_vec(); + self.forward_scratch = r2c.make_scratch_vec(); + self.inverse_scratch = c2r.make_scratch_vec(); + self.time_domain = c2r.make_output_vec(); + self.plans = Some(plans); + self.fft_size = fft_size; + } - r2c.process_with_scratch(&mut padded_signal, &mut signal_spectrum, &mut forward_scratch) - .expect("forward FFT buffers are sized by the plan"); - r2c.process_with_scratch(&mut padded_template, &mut template_spectrum, &mut forward_scratch) - .expect("forward FFT buffers are sized by the plan"); + self.padded_signal.clear(); + self.padded_signal.resize(fft_size, 0.0); + self.padded_signal[..self.signal.len()].copy_from_slice(&self.signal); - for (s, t) in signal_spectrum.iter_mut().zip(template_spectrum.iter()) { - *s *= *t; + let (r2c, _) = self.plans.as_ref().expect("plans were just set"); + r2c.process_with_scratch(&mut self.padded_signal, &mut self.signal_spectrum, &mut self.forward_scratch) + .expect("forward FFT buffers are sized by the plan"); + self.spectrum_loaded = true; } - let mut result = c2r.make_output_vec(); - let mut inverse_scratch = c2r.make_scratch_vec(); - c2r.process_with_scratch(&mut signal_spectrum, &mut result, &mut inverse_scratch) - .expect("inverse FFT buffers are sized by the plan"); + /// Full cross-correlation of the loaded signal with `template`, written + /// to `output` (length `signal.len() + template.len() - 1`). + /// + /// Matches `scipy.signal.correlate(signal, template, mode="full")`: + /// output index `k` is the lag where `template[template.len() - 1]` + /// aligns with `signal[k]`. `output` is empty if either is empty. + pub fn correlate(&mut self, template: &mut CorrelationTemplate, output: &mut Vec) { + output.clear(); + if self.signal.is_empty() || template.is_empty() { + return; + } - result.truncate(output_len); - let normalization = 1.0 / fft_size as f32; - result.iter_mut().for_each(|x| *x *= normalization); - result + // Covers the full correlation without wrap-around. + let output_len = self.signal.len() + template.len() - 1; + self.transform_signal(output_len.next_power_of_two()); + + let template_spectrum = template.spectrum(self.fft_size); + for ((p, s), t) in self.product.iter_mut().zip(&self.signal_spectrum).zip(template_spectrum) { + *p = *s * *t; + } + + let (_, c2r) = self.plans.as_ref().expect("the signal was just transformed"); + c2r.process_with_scratch(&mut self.product, &mut self.time_domain, &mut self.inverse_scratch) + .expect("inverse FFT buffers are sized by the plan"); + + let normalization = 1.0 / self.fft_size as f32; + output.extend(self.time_domain[..output_len].iter().map(|x| x * normalization)); + } +} + +/// Cross-correlate two real 1-D signals using FFT, returning the full +/// correlation (length `signal.len() + template.len() - 1`). +/// +/// One-shot convenience over [`CorrelationWorkspace`]; see +/// [`CorrelationWorkspace::correlate`] for the semantics. +pub fn fft_correlate_full(signal: &[f32], template: &[f32]) -> Vec { + let mut workspace = CorrelationWorkspace::new(); + let mut template = CorrelationTemplate::new(template); + let mut output = Vec::new(); + workspace.load_signal(signal); + workspace.correlate(&mut template, &mut output); + output } #[cfg(test)] @@ -116,6 +233,8 @@ mod tests { fn test_empty_inputs() { assert_eq!(fft_correlate_full(&[], &[1.0]), Vec::::new()); assert_eq!(fft_correlate_full(&[1.0], &[]), Vec::::new()); + assert_eq!(fft_correlate_full(&[1.0, 2.0], &[]), Vec::::new()); + assert_eq!(fft_correlate_full(&[1.0, 2.0, 3.0], &[]), Vec::::new()); } #[test] @@ -130,6 +249,175 @@ mod tests { } } + #[test] + fn test_workspace_reuse_matches_one_shot() { + let signal_a: Vec = (0..50).map(|i| ((i * 13 % 17) as f32 - 8.0) / 8.0).collect(); + let signal_b: Vec = (0..70).map(|i| ((i * 3 % 19) as f32 - 9.0) / 9.0).collect(); + let template_a: Vec = (0..7).map(|i| ((i * 5 % 7) as f32 - 3.0) / 3.0).collect(); + let template_b: Vec = (0..12).map(|i| ((i * 11 % 13) as f32 - 6.0) / 6.0).collect(); + + let mut workspace = CorrelationWorkspace::new(); + let mut cached_a = CorrelationTemplate::new(&template_a); + let mut cached_b = CorrelationTemplate::new(&template_b); + let mut output = Vec::new(); + + // Two templates against one signal, then a different signal size. + for signal in [&signal_a, &signal_b, &signal_a] { + workspace.load_signal(signal); + for (cached, template) in [(&mut cached_a, &template_a), (&mut cached_b, &template_b)] { + workspace.correlate(cached, &mut output); + let naive = naive_full_correlation(signal, template); + assert_eq!(output.len(), naive.len()); + for (a, b) in output.iter().zip(naive.iter()) { + assert!((a - b).abs() < 1e-4, "{a} != {b}"); + } + } + } + assert_eq!(cached_a.spectra.len(), 2, "one cached spectrum per FFT size"); + } + + #[test] + fn test_workspace_empty_inputs() { + let mut workspace = CorrelationWorkspace::new(); + let mut template = CorrelationTemplate::new(&[1.0, 2.0]); + let mut output = vec![1.0]; + workspace.load_signal(&[]); + workspace.correlate(&mut template, &mut output); + assert!(output.is_empty()); + + let mut empty = CorrelationTemplate::new(&[]); + workspace.load_signal(&[1.0, 2.0, 3.0]); + workspace.correlate(&mut empty, &mut output); + assert!(output.is_empty()); + } + + fn assert_close(actual: &[f32], expected: &[f32]) { + assert_eq!(actual.len(), expected.len()); + for (a, b) in actual.iter().zip(expected) { + assert!((a - b).abs() < 1e-4, "{a} != {b}"); + } + } + + #[test] + fn test_template_len_and_is_empty() { + let template = CorrelationTemplate::new(&[1.0, 2.0, 3.0]); + assert_eq!(template.len(), 3); + assert!(!template.is_empty()); + assert_eq!(template.reversed, vec![3.0, 2.0, 1.0]); + + let empty = CorrelationTemplate::new(&[]); + assert_eq!(empty.len(), 0); + assert!(empty.is_empty()); + } + + #[test] + fn test_cached_spectrum_gives_identical_output() { + let signal: Vec = (0..90).map(|i| ((i * 29 % 31) as f32 - 15.0) / 15.0).collect(); + let template: Vec = (0..11).map(|i| ((i * 7 % 13) as f32 - 6.0) / 6.0).collect(); + + let mut workspace = CorrelationWorkspace::new(); + let mut cached = CorrelationTemplate::new(&template); + let mut first = Vec::new(); + let mut second = Vec::new(); + + workspace.load_signal(&signal); + workspace.correlate(&mut cached, &mut first); + assert_eq!(cached.spectra.len(), 1); + // Same signal again: the cached spectrum is used (no new entry) and + // the result is bit-identical to the first pass and to the one-shot. + workspace.load_signal(&signal); + workspace.correlate(&mut cached, &mut second); + assert_eq!(cached.spectra.len(), 1); + assert_eq!(first, second); + assert_eq!(first, fft_correlate_full(&signal, &template)); + } + + #[test] + fn test_template_spectrum_cache_evicts_oldest() { + let template: Vec = (0..5).map(|i| ((i * 3 % 5) as f32 - 2.0) / 2.0).collect(); + let mut cached = CorrelationTemplate::new(&template); + let mut workspace = CorrelationWorkspace::new(); + let mut output = Vec::new(); + + // Six signal lengths, each needing a different power-of-two FFT. + let lengths = [4, 12, 28, 60, 124, 252]; + let mut fft_sizes = Vec::new(); + for &len in &lengths { + let signal: Vec = (0..len).map(|i| ((i * 17 % 23) as f32 - 11.0) / 11.0).collect(); + workspace.load_signal(&signal); + workspace.correlate(&mut cached, &mut output); + assert_close(&output, &naive_full_correlation(&signal, &template)); + fft_sizes.push(workspace.fft_size); + assert!(cached.spectra.len() <= TEMPLATE_SPECTRUM_CACHE_CAPACITY); + } + assert_eq!(fft_sizes, vec![8, 16, 32, 64, 128, 256]); + + // The oldest two sizes were evicted, the newest four kept in order. + let kept: Vec = cached.spectra.iter().map(|(size, _)| *size).collect(); + assert_eq!(kept, vec![32, 64, 128, 256]); + + // An evicted size is recomputed and still correct. + let signal: Vec = (0..4).map(|i| i as f32 - 1.5).collect(); + workspace.load_signal(&signal); + workspace.correlate(&mut cached, &mut output); + assert_close(&output, &naive_full_correlation(&signal, &template)); + let kept: Vec = cached.spectra.iter().map(|(size, _)| *size).collect(); + assert_eq!(kept, vec![64, 128, 256, 8]); + } + + #[test] + fn test_fft_size_is_independent_of_other_templates() { + let signal: Vec = (0..40).map(|i| ((i * 19 % 29) as f32 - 14.0) / 14.0).collect(); + let short: Vec = (0..6).map(|i| ((i * 5 % 7) as f32 - 3.0) / 3.0).collect(); + let long: Vec = (0..30).map(|i| ((i * 11 % 17) as f32 - 8.0) / 8.0).collect(); + let expected_short = fft_correlate_full(&signal, &short); + let expected_long = fft_correlate_full(&signal, &long); + + let mut workspace = CorrelationWorkspace::new(); + let mut output = Vec::new(); + workspace.load_signal(&signal); + // Each template gets the FFT size it needs on its own, in any order, + // so the output is bit-identical to the one-shot correlation. + let mut fft_sizes = Vec::new(); + for (template, expected) in [(&long, &expected_long), (&short, &expected_short), (&long, &expected_long)] { + workspace.correlate(&mut CorrelationTemplate::new(template), &mut output); + fft_sizes.push(workspace.fft_size); + assert_eq!(&output, expected); + assert_close(&output, &naive_full_correlation(&signal, template)); + } + assert_eq!(fft_sizes, vec![128, 64, 128]); + } + + #[test] + fn test_template_longer_than_signal_matches_naive() { + let signal = [1.0_f32, -2.0, 0.5]; + let template: Vec = (0..9).map(|i| ((i * 4 % 9) as f32 - 4.0) / 4.0).collect(); + let result = fft_correlate_full(&signal, &template); + assert_eq!(result.len(), 11); + assert_close(&result, &naive_full_correlation(&signal, &template)); + } + + #[test] + fn test_exact_values_match_scipy() { + // scipy.signal.correlate([1, 2, 3, 4], [1, 0.5], mode="full") + // = [0.5, 2.0, 3.5, 5.0, 4.0] + let result = fft_correlate_full(&[1.0, 2.0, 3.0, 4.0], &[1.0, 0.5]); + assert_close(&result, &[0.5, 2.0, 3.5, 5.0, 4.0]); + // Single-sample template scales the signal. + let result = fft_correlate_full(&[1.0, -2.0, 3.0], &[2.0]); + assert_close(&result, &[2.0, -4.0, 6.0]); + } + + #[test] + fn test_output_buffer_is_replaced_not_appended() { + let mut workspace = CorrelationWorkspace::new(); + let mut template = CorrelationTemplate::new(&[1.0, 0.5]); + let mut output = vec![9.0; 20]; + workspace.load_signal(&[1.0, 2.0, 3.0, 4.0]); + workspace.correlate(&mut template, &mut output); + assert_close(&output, &[0.5, 2.0, 3.5, 5.0, 4.0]); + } + #[test] fn test_autocorrelation_peak_is_centered() { let clip = [0.5_f32, -1.0, 0.25, 0.75]; diff --git a/crates/core/src/dsp/loudness.rs b/crates/core/src/dsp/loudness.rs index 0569bbc..ff57788 100644 --- a/crates/core/src/dsp/loudness.rs +++ b/crates/core/src/dsp/loudness.rs @@ -207,17 +207,20 @@ pub fn integrated_loudness(data: &[f32], sample_rate: u32, block_size: f64) -> f LUFS_OFFSET + 10.0 * z_avg_final.log10() } -/// Normalize audio to a target loudness in dB LUFS with hard clipping. +/// Normalize audio in place to a target loudness in dB LUFS with hard +/// clipping. /// /// Applies the gain needed to shift from `current_lufs` to `target_lufs`, -/// then hard-clips the output to [-1.0, 1.0]. -pub fn loudness_normalize(data: &[f32], current_lufs: f64, target_lufs: f64) -> Vec { +/// then hard-clips the output to [-1.0, 1.0]. Silence (`current_lufs` of +/// `-inf`, hence an infinite gain) leaves zero samples at zero instead of +/// the NaN that `0 * inf` would give. +pub fn loudness_normalize(data: &mut [f32], current_lufs: f64, target_lufs: f64) { let delta = target_lufs - current_lufs; let gain = 10.0_f64.powf(delta / 20.0); - data.iter() - .map(|&x| ((x as f64) * gain).clamp(-1.0, 1.0) as f32) - .collect() + for x in data.iter_mut().filter(|x| **x != 0.0) { + *x = ((*x as f64) * gain).clamp(-1.0, 1.0) as f32; + } } #[cfg(test)] @@ -300,19 +303,53 @@ mod tests { #[test] fn test_loudness_normalize_clips() { - let data = [0.5_f32, -0.5, 0.8, -0.8]; + let mut out = [0.5_f32, -0.5, 0.8, -0.8]; // Apply huge gain (+40 dB) to force clipping. - let out = loudness_normalize(&data, -60.0, -20.0); + loudness_normalize(&mut out, -60.0, -20.0); for &v in &out { assert!((-1.0..=1.0).contains(&v), "value {v} exceeds [-1, 1]"); } } + #[test] + fn test_loudness_normalize_in_place_exact() { + // +20 dB is exactly a gain of 10; binary fractions keep the + // products exact, and -0.125 * 10 is hard-clipped. + let mut up = [0.0625_f32, -0.125, 0.03125, 0.0]; + loudness_normalize(&mut up, -36.0, -16.0); + assert_eq!(up, [0.625, -1.0, 0.3125, 0.0]); + + // -20 dB attenuates by 10. + let mut down = [0.5_f32, -1.0, 0.25]; + loudness_normalize(&mut down, 4.0, -16.0); + for (actual, expected) in down.iter().zip([0.05_f32, -0.1, 0.025]) { + assert!((actual - expected).abs() < 1e-7, "{actual} != {expected}"); + } + + // Zero gain change leaves every sample untouched, clipping aside. + let mut same = [0.3_f32, -0.7, 1.5, -2.0]; + loudness_normalize(&mut same, -16.0, -16.0); + assert_eq!(same, [0.3, -0.7, 1.0, -1.0]); + + // Silence measures -inf LUFS, so the gain is infinite: zeros stay + // zero (not NaN) and anything else is hard-clipped to full scale. + let mut silent = [0.0_f32, 0.0, -0.0]; + loudness_normalize(&mut silent, f64::NEG_INFINITY, -16.0); + assert_eq!(silent, [0.0, 0.0, 0.0]); + let mut sub_gate = [0.0_f32, 1e-6, -1e-6]; + loudness_normalize(&mut sub_gate, f64::NEG_INFINITY, -16.0); + assert_eq!(sub_gate, [0.0, 1.0, -1.0]); + + let mut empty: [f32; 0] = []; + loudness_normalize(&mut empty, -20.0, -16.0); + assert!(empty.is_empty()); + } + #[test] fn test_loudness_normalize_gain() { - let data = [0.1_f32, -0.1]; + let mut out = [0.1_f32, -0.1]; // +6 dB gain ≈ 2x. - let out = loudness_normalize(&data, -22.0, -16.0); + loudness_normalize(&mut out, -22.0, -16.0); let expected_gain = 10.0_f64.powf(6.0 / 20.0); // ~1.995 assert!((out[0] as f64 - 0.1 * expected_gain).abs() < 1e-4); } diff --git a/crates/core/src/dsp/mod.rs b/crates/core/src/dsp/mod.rs index 34d2e00..266d0e8 100644 --- a/crates/core/src/dsp/mod.rs +++ b/crates/core/src/dsp/mod.rs @@ -8,7 +8,7 @@ pub mod peaks; pub mod resample; pub mod spectrum; -pub use correlate::fft_correlate_full; +pub use correlate::{fft_correlate_full, CorrelationTemplate, CorrelationWorkspace}; pub use loudness::{integrated_loudness, loudness_normalize}; pub use pearson::pearson_correlation_1d; pub use peaks::{find_peaks_1d, FindPeaksOptions}; diff --git a/crates/core/src/dsp/peaks.rs b/crates/core/src/dsp/peaks.rs index 37d465d..2d93086 100644 --- a/crates/core/src/dsp/peaks.rs +++ b/crates/core/src/dsp/peaks.rs @@ -14,11 +14,10 @@ pub struct FindPeaksOptions { /// /// Returns a sorted vector of peak indices. pub fn find_peaks_1d(data: &[f32], options: &FindPeaksOptions) -> Vec { - let mut peaks = local_maxima_1d(data); - - if let Some(min_height) = options.height { - filter_by_height(data, &mut peaks, min_height); - } + // The height condition is applied while scanning: it depends only on the + // peak's own value, so the result is the same as filtering afterwards, + // and far fewer candidates are collected. + let mut peaks = local_maxima_1d(data, options.height); if let Some(min_distance) = options.distance { filter_by_distance(data, &mut peaks, min_distance); @@ -31,43 +30,58 @@ pub fn find_peaks_1d(data: &[f32], options: &FindPeaksOptions) -> Vec { peaks } -/// Detect all local maxima in `data`. +/// Samples per block of the local-maxima scan; the candidate test for a +/// block is written as straight-line slice arithmetic so it vectorizes. +const SCAN_BLOCK: usize = 64; + +/// Detect all local maxima in `data` whose value is at least `min_height`. /// /// A sample is a local maximum when it is strictly greater than both its /// immediate neighbours. For plateaus (runs of identical values that are /// higher than the values on both sides) the midpoint index (rounded down) /// is returned, matching scipy's behaviour. -fn local_maxima_1d(data: &[f32]) -> Vec { +fn local_maxima_1d(data: &[f32], min_height: Option) -> Vec { let n = data.len(); if n < 3 { return vec![]; } + let min_height = min_height.unwrap_or(f32::NEG_INFINITY); let mut peaks = Vec::new(); - let mut i = 1; - while i < n - 1 { - if data[i - 1] < data[i] { - let left_edge = i; - // Advance through equal values (plateau). - while i + 1 < n && data[i] == data[i + 1] { - i += 1; - } - let right_edge = i; - // Confirm the right side drops. - if i + 1 < n && data[i] > data[i + 1] { - peaks.push((left_edge + right_edge) / 2); + let mut candidates = [false; SCAN_BLOCK]; + let mut start = 1; + while start < n - 1 { + let end = (start + SCAN_BLOCK).min(n - 1); + let block = &mut candidates[..end - start]; + let previous = &data[start - 1..end - 1]; + let current = &data[start..end]; + let next = &data[start + 1..end + 1]; + + // A candidate rises from the left, does not keep rising (strict peak + // or plateau start) and is tall enough. A NaN neighbour fails the + // comparison, just as it would fail the plateau check below. + for (((flag, &p), &c), &nx) in block.iter_mut().zip(previous).zip(current).zip(next) { + *flag = p < c && c >= nx && c >= min_height; + } + if block.iter().any(|&flag| flag) { + for (offset, _) in block.iter().enumerate().filter(|(_, &flag)| flag) { + let left_edge = start + offset; + // Advance through equal values (plateau). + let mut right_edge = left_edge; + while right_edge + 1 < n && data[right_edge] == data[right_edge + 1] { + right_edge += 1; + } + // Confirm the right side drops. + if right_edge + 1 < n && data[right_edge] > data[right_edge + 1] { + peaks.push((left_edge + right_edge) / 2); + } } } - i += 1; + start = end; } peaks } -/// Keep only peaks whose value is at least `min_height`. -fn filter_by_height(data: &[f32], peaks: &mut Vec, min_height: f32) { - peaks.retain(|&idx| data[idx] >= min_height); -} - /// Keep only the tallest peaks when multiple peaks fall within `min_distance` /// of each other. Matches scipy's greedy tallest-first strategy. fn filter_by_distance(data: &[f32], peaks: &mut Vec, min_distance: usize) { @@ -287,37 +301,206 @@ mod tests { #[test] fn test_local_maxima_simple() { let data = [0.0, 1.0, 0.0, 2.0, 0.0]; - assert_eq!(local_maxima_1d(&data), vec![1, 3]); + assert_eq!(local_maxima_1d(&data, None), vec![1, 3]); } #[test] fn test_local_maxima_empty_and_short() { - assert_eq!(local_maxima_1d(&[]), Vec::::new()); - assert_eq!(local_maxima_1d(&[1.0]), Vec::::new()); - assert_eq!(local_maxima_1d(&[1.0, 2.0]), Vec::::new()); + assert_eq!(local_maxima_1d(&[], None), Vec::::new()); + assert_eq!(local_maxima_1d(&[1.0], None), Vec::::new()); + assert_eq!(local_maxima_1d(&[1.0, 2.0], None), Vec::::new()); } #[test] fn test_local_maxima_plateau_even() { // Plateau [1,1] spans indices 1..2 → midpoint = 1 let data = [0.0, 1.0, 1.0, 0.0]; - assert_eq!(local_maxima_1d(&data), vec![1]); + assert_eq!(local_maxima_1d(&data, None), vec![1]); } #[test] fn test_local_maxima_plateau_odd() { // Plateau [1,1,1] spans indices 1..3 → midpoint = 2 let data = [0.0, 1.0, 1.0, 1.0, 0.0]; - assert_eq!(local_maxima_1d(&data), vec![2]); + assert_eq!(local_maxima_1d(&data, None), vec![2]); } #[test] fn test_local_maxima_monotonic() { let ascending = [1.0, 2.0, 3.0, 4.0, 5.0]; - assert_eq!(local_maxima_1d(&ascending), Vec::::new()); + assert_eq!(local_maxima_1d(&ascending, None), Vec::::new()); let descending = [5.0, 4.0, 3.0, 2.0, 1.0]; - assert_eq!(local_maxima_1d(&descending), Vec::::new()); + assert_eq!(local_maxima_1d(&descending, None), Vec::::new()); + } + + #[test] + fn test_local_maxima_height_in_scan() { + let data = [0.0, 1.0, 0.0, 2.0, 2.0, 0.0, 3.0, 0.0]; + assert_eq!(local_maxima_1d(&data, None), vec![1, 3, 6]); + assert_eq!(local_maxima_1d(&data, Some(2.0)), vec![3, 6]); + assert_eq!(local_maxima_1d(&data, Some(2.5)), vec![6]); + assert_eq!(local_maxima_1d(&data, Some(4.0)), Vec::::new()); + } + + #[test] + fn test_local_maxima_across_scan_blocks() { + // Peaks at the last and first positions of a block, a plateau that + // starts in one block and ends in the next, and a NaN neighbour. + let mut data = vec![0.0_f32; 3 * SCAN_BLOCK]; + data[SCAN_BLOCK] = 1.0; // start of block 2 (scan starts at index 1) + data[SCAN_BLOCK + 1] = 1.0; + data[2 * SCAN_BLOCK] = 2.0; // last position of block 2 + data[2 * SCAN_BLOCK + 1] = 2.0; // plateau continues into block 3 + data[2 * SCAN_BLOCK + 2] = 2.0; + data[2 * SCAN_BLOCK + 10] = 5.0; + data[2 * SCAN_BLOCK + 11] = f32::NAN; + data[2 * SCAN_BLOCK + 20] = 4.0; + assert_eq!( + local_maxima_1d(&data, None), + vec![SCAN_BLOCK, 2 * SCAN_BLOCK + 1, 2 * SCAN_BLOCK + 20] + ); + assert_eq!(local_maxima_1d(&data, Some(1.5)), vec![2 * SCAN_BLOCK + 1, 2 * SCAN_BLOCK + 20]); + } + + /// Plain one-pass scan with the height applied afterwards: the + /// scipy `_local_maxima_1d` algorithm the block-wise scan must match. + fn reference_local_maxima(data: &[f32], min_height: Option) -> Vec { + let mut peaks = Vec::new(); + let mut i = 1; + while i + 1 < data.len() { + if data[i - 1] < data[i] { + let left_edge = i; + while i + 1 < data.len() && data[i] == data[i + 1] { + i += 1; + } + if i + 1 < data.len() && data[i] > data[i + 1] { + peaks.push((left_edge + i) / 2); + } + } + i += 1; + } + if let Some(min_height) = min_height { + peaks.retain(|&idx| data[idx] >= min_height); + } + peaks + } + + /// Pseudo-random values with many repeats (so plateaus are common). + fn repetitive_data(n: usize) -> Vec { + (0..n).map(|i| ((i * 7919 % 23) as f32 - 11.0) / 4.0).collect() + } + + #[test] + fn test_local_maxima_matches_reference_scan() { + let data = repetitive_data(1000); + let expected = reference_local_maxima(&data, None); + assert!(!expected.is_empty()); + assert_eq!(local_maxima_1d(&data, None), expected); + let tall = reference_local_maxima(&data, Some(2.0)); + assert_ne!(tall.len(), expected.len()); + assert_eq!(local_maxima_1d(&data, Some(2.0)), tall); + } + + #[test] + fn test_local_maxima_matches_reference_for_every_length() { + // Every length up to a few blocks, so the last block takes every + // possible partial size and data ends at every block offset. + for n in 0..=3 * SCAN_BLOCK + 3 { + let data = repetitive_data(n); + for min_height in [None, Some(-1.0), Some(0.0), Some(2.75), Some(3.0)] { + assert_eq!( + local_maxima_1d(&data, min_height), + reference_local_maxima(&data, min_height), + "length {n}, min_height {min_height:?}" + ); + } + } + } + + #[test] + fn test_local_maxima_alternating_hits_every_block_position() { + // 0,1,0,1,...: a peak at every odd index, so each position of + // each block (first, last and interior) is a peak in one of the + // two phases. + let n = 3 * SCAN_BLOCK + 2; + let odd_peaks: Vec = (0..n).map(|i| (i % 2) as f32).collect(); + let expected: Vec = (1..n - 1).filter(|i| i % 2 == 1).collect(); + assert_eq!(local_maxima_1d(&odd_peaks, None), expected); + + let even_peaks: Vec = (0..n).map(|i| ((i + 1) % 2) as f32).collect(); + let expected: Vec = (1..n - 1).filter(|i| i % 2 == 0).collect(); + assert_eq!(local_maxima_1d(&even_peaks, None), expected); + assert_eq!(local_maxima_1d(&even_peaks, Some(1.5)), Vec::::new()); + } + + #[test] + fn test_local_maxima_plateau_spanning_whole_block() { + // A plateau longer than a block that starts in block 1 and drops in + // block 3 is reported once, at its midpoint. + let mut data = vec![0.0_f32; 4 * SCAN_BLOCK]; + let left = SCAN_BLOCK / 2; + let right = left + SCAN_BLOCK + 10; + data[left..=right].fill(1.0); + assert_eq!(local_maxima_1d(&data, None), vec![(left + right) / 2]); + assert_eq!(local_maxima_1d(&data, Some(1.0)), vec![(left + right) / 2]); + assert_eq!(local_maxima_1d(&data, Some(1.1)), Vec::::new()); + + // A plateau that runs to the end of the data never drops: no peak. + data[right + 1..].fill(1.0); + assert_eq!(local_maxima_1d(&data, None), Vec::::new()); + } + + #[test] + fn test_local_maxima_nan_handling_matches_reference() { + // NaN as the peak value, as the left neighbour and as the right + // neighbour, plus NaN inside a plateau. None of them is a peak and + // none hides a real peak elsewhere. + let mut data = vec![0.0_f32; 40]; + data[2] = f32::NAN; // NaN "peak" + data[6] = f32::NAN; // NaN left of a rise + data[7] = 1.0; + data[12] = 1.0; + data[13] = f32::NAN; // NaN right of a peak + data[20] = 1.0; // plateau broken by NaN + data[21] = f32::NAN; + data[22] = 1.0; + data[30] = 2.0; // ordinary peak + let expected = reference_local_maxima(&data, None); + assert_eq!(expected, vec![30]); + assert_eq!(local_maxima_1d(&data, None), expected); + assert_eq!(local_maxima_1d(&data, Some(1.0)), vec![30]); + assert_eq!(local_maxima_1d(&data, Some(f32::NAN)), Vec::::new()); + } + + #[test] + fn test_find_peaks_height_matches_filtering_afterwards() { + // find_peaks with height must equal find_peaks without height + // followed by a height filter, including with distance and + // prominence applied after it. + let data = repetitive_data(3 * SCAN_BLOCK + 7); + for height in [None, Some(0.5), Some(2.0)] { + for distance in [None, Some(3)] { + for prominence in [None, Some(1.0)] { + let with_height = find_peaks_1d(&data, &FindPeaksOptions { height, distance, prominence }); + let mut unfiltered = + find_peaks_1d(&data, &FindPeaksOptions { height: None, distance: None, prominence: None }); + if let Some(min_height) = height { + unfiltered.retain(|&idx| data[idx] >= min_height); + } + if let Some(min_distance) = distance { + filter_by_distance(&data, &mut unfiltered, min_distance); + } + if let Some(min_prominence) = prominence { + filter_by_prominence(&data, &mut unfiltered, min_prominence); + } + assert_eq!( + with_height, unfiltered, + "height {height:?} distance {distance:?} prominence {prominence:?}" + ); + } + } + } } // ── height filter ──────────────────────────────────────────────── diff --git a/crates/core/tests/sliding_window.rs b/crates/core/tests/sliding_window.rs index 079bc8c..ef4b2d0 100644 --- a/crates/core/tests/sliding_window.rs +++ b/crates/core/tests/sliding_window.rs @@ -10,7 +10,7 @@ use audio_pattern_detector_core::stream::WavFileSource; use audio_pattern_detector_core::{ match_pattern, AudioClip, AudioPatternDetector, AudioStream, DetectorOptions, MatchOptions, PeakTimes, }; -use common::{clip_from_samples, concat, insert_at, silence, sine_tone, stream_from_samples, SR}; +use common::{clip_from_samples, concat, insert_at, silence, sine_tone, stream_from_samples, Rng, SR}; const RTHK_BEEP_PATTERN: &str = "../../sample_audios/clips/rthk_beep.apd.toml"; const RTHK_BEEP_AUDIO: &str = "../../sample_audios/rthk_section_with_beep.wav"; @@ -238,6 +238,176 @@ mod sliding_window_boundary { } } +/// Each clip's section is loudness-normalized on its own lookback, so one +/// clip's detections never depend on which other clips are loaded. +mod per_clip_normalization { + use super::*; + + const UNRELATED_NAME: &str = "unrelated_long"; + + /// Loud unrelated audio at 6.5-8.5s, then a quiet beep at 11s. With 10s + /// chunks the beep's own 1s lookback section (9-20s) holds only the + /// beep, while a 4s lookback (6-20s) would also take in the loud audio. + fn quiet_beep_after_loud_audio() -> Vec { + let mut audio = silence(30.0, SR); + let loud: Vec = sine_tone(300.0, 2.0, SR).iter().map(|v| v * 0.9).collect(); + insert_at(&mut audio, (6.5 * SR as f64) as usize, &loud); + let quiet: Vec = beep_pattern().audio.iter().map(|v| v * 0.02).collect(); + insert_at(&mut audio, (11.0 * SR as f64) as usize, &quiet); + audio + } + + fn beep_detections(clips: Vec) -> Vec { + let detector = new_detector(clips, Some(10)).unwrap(); + let mut stream = stream_from_samples("test_audio", &quiet_beep_after_loud_audio()); + let (peak_times, _) = detector.find_clip_in_audio(&mut stream, None, true).unwrap(); + detections(&peak_times.unwrap(), BEEP_NAME).to_vec() + } + + #[test] + fn unrelated_longer_clip_does_not_change_detections() { + let alone = beep_detections(vec![beep_pattern()]); + assert_eq!(alone.len(), 1, "quiet beep alone: {alone:?}"); + assert!((alone[0] - 11.0).abs() < 0.05, "quiet beep alone at {alone:?}"); + + let unrelated = tone_clip(UNRELATED_NAME, 2000.0, 3.5); + let with_unrelated = beep_detections(vec![beep_pattern(), unrelated]); + assert_eq!(with_unrelated, alone, "detections changed when an unrelated clip was added"); + } +} + +/// Clips with the same sliding window share one loudness-normalized section +/// and one forward FFT (sized for the longest clip of the group). Sharing +/// must not change any clip's detections compared with running it alone. +mod shared_sections { + use super::*; + + const MID_NAME: &str = "mid_tone"; + const MID_DURATION: f64 = 0.6; + const MID_FREQUENCY: f64 = 2000.0; + const LOW_NAME: &str = "low_long"; + const LOW_FREQUENCY: f64 = 500.0; + + /// 0.6s 2 kHz tone: sliding_window 1 like the beep, but longer, so + /// the shared FFT is sized for it rather than for the beep. + fn mid_pattern() -> AudioClip { + tone_clip(MID_NAME, MID_FREQUENCY, MID_DURATION) + } + + /// 2.5s 500 Hz tone: sliding_window 3, its own group. + fn low_pattern() -> AudioClip { + tone_clip(LOW_NAME, LOW_FREQUENCY, LONG_BEEP_DURATION) + } + + // 10s chunks: positions in the first chunk, in a later chunk, and + // straddling a chunk boundary (so the lookback matters). + const BEEP_POSITIONS: &[f64] = &[3.0, 9.9, 15.0]; + const MID_POSITIONS: &[f64] = &[5.0, 12.0, 19.8]; + const LOW_POSITIONS: &[f64] = &[0.5, 24.0]; + const AUDIO_SECONDS: f64 = 30.0; + const SECONDS_PER_CHUNK: u32 = 10; + + fn mixed_audio() -> Vec { + let mut audio = silence(AUDIO_SECONDS, SR); + for (pattern, positions) in + [(beep_pattern(), BEEP_POSITIONS), (mid_pattern(), MID_POSITIONS), (low_pattern(), LOW_POSITIONS)] + { + for &pos in positions { + insert_at(&mut audio, (pos * SR as f64) as usize, &pattern.audio); + } + } + audio + } + + fn run(clips: Vec) -> PeakTimes { + let detector = new_detector(clips, Some(SECONDS_PER_CHUNK)).unwrap(); + let mut stream = stream_from_samples("test_audio", &mixed_audio()); + let (peak_times, _) = detector.find_clip_in_audio(&mut stream, None, true).unwrap(); + peak_times.expect("results are accumulated") + } + + fn assert_found_at(times: &[f64], positions: &[f64], name: &str) { + assert_eq!(times.len(), positions.len(), "{name}: {times:?} vs {positions:?}"); + for (actual, expected) in times.iter().zip(positions) { + assert!((actual - expected).abs() < 0.05, "{name}: {actual} vs {expected} (all: {times:?})"); + } + } + + #[test] + fn each_clip_alone_finds_every_occurrence() { + assert_found_at(detections(&run(vec![beep_pattern()]), BEEP_NAME), BEEP_POSITIONS, BEEP_NAME); + assert_found_at(detections(&run(vec![mid_pattern()]), MID_NAME), MID_POSITIONS, MID_NAME); + assert_found_at(detections(&run(vec![low_pattern()]), LOW_NAME), LOW_POSITIONS, LOW_NAME); + } + + #[test] + fn clips_sharing_a_section_match_their_solo_runs_exactly() { + let beep_alone = detections(&run(vec![beep_pattern()]), BEEP_NAME).to_vec(); + let mid_alone = detections(&run(vec![mid_pattern()]), MID_NAME).to_vec(); + + // Same group, shorter clip first (FFT sized for the second clip). + let together = run(vec![beep_pattern(), mid_pattern()]); + assert_eq!(detections(&together, BEEP_NAME), beep_alone); + assert_eq!(detections(&together, MID_NAME), mid_alone); + + // Same group, longer clip first. + let together = run(vec![mid_pattern(), beep_pattern()]); + assert_eq!(detections(&together, BEEP_NAME), beep_alone); + assert_eq!(detections(&together, MID_NAME), mid_alone); + } + + #[test] + fn clips_in_different_groups_match_their_solo_runs_exactly() { + let beep_alone = detections(&run(vec![beep_pattern()]), BEEP_NAME).to_vec(); + let mid_alone = detections(&run(vec![mid_pattern()]), MID_NAME).to_vec(); + let low_alone = detections(&run(vec![low_pattern()]), LOW_NAME).to_vec(); + + // Groups in both orders of first appearance. + for clips in [ + vec![beep_pattern(), low_pattern(), mid_pattern()], + vec![low_pattern(), mid_pattern(), beep_pattern()], + ] { + let together = run(clips); + assert_eq!(detections(&together, BEEP_NAME), beep_alone); + assert_eq!(detections(&together, MID_NAME), mid_alone); + assert_eq!(detections(&together, LOW_NAME), low_alone); + } + } + + #[test] + fn duplicate_clips_get_identical_detections() { + let mut copy = beep_pattern(); + copy.name = "beep_copy".to_string(); + let together = run(vec![beep_pattern(), copy]); + let beep = detections(&together, BEEP_NAME); + assert_found_at(beep, BEEP_POSITIONS, BEEP_NAME); + assert_eq!(detections(&together, "beep_copy"), beep); + } + + #[test] + fn callback_order_is_by_time_within_a_chunk_across_groups() { + use std::sync::{Arc, Mutex}; + let clips = vec![beep_pattern(), low_pattern(), mid_pattern()]; + let detector = new_detector(clips, Some(SECONDS_PER_CHUNK)).unwrap(); + let mut stream = stream_from_samples("test_audio", &mixed_audio()); + let events: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sink = Arc::clone(&events); + let mut callback = move |name: &str, time: f64| sink.lock().unwrap().push((time, name.to_string())); + let (peak_times, _) = detector.find_clip_in_audio(&mut stream, Some(&mut callback), true).unwrap(); + let peak_times = peak_times.unwrap(); + + let events = events.lock().unwrap(); + let times: Vec = events.iter().map(|(t, _)| *t).collect(); + assert!(times.windows(2).all(|w| w[0] <= w[1]), "callback not in time order: {events:?}"); + + // The callback saw exactly the accumulated detections of every clip. + for name in [BEEP_NAME, MID_NAME, LOW_NAME] { + let from_callback: Vec = events.iter().filter(|(_, n)| n == name).map(|(t, _)| *t).collect(); + assert_eq!(from_callback, detections(&peak_times, name), "{name}"); + } + } +} + /// Integration tests using real audio patterns for sliding window behaviour. mod sliding_window_with_real_patterns { use super::*; @@ -769,3 +939,80 @@ fn test_overlap_detections_are_reported_once() { expected.iter().map(|&t| (SHORT_BEEP_NAME.to_string(), t)).collect(); assert_eq!(events, expected_events); } + +const TIE_OTHER_NAME: &str = "tie_other"; +const TIE_PREFIX_NAME: &str = "tie_prefix"; +const TIE_FULL_NAME: &str = "tie_full"; +const TIE_PREFIX_DURATION: f64 = 0.5; +const TIE_FULL_DURATION: f64 = 1.5; +const TIE_START: f64 = 2.0; +const TIE_EXPECTED_TIME: f64 = 1.9998749999999998; + +// Clips are grouped by sliding window for matching, but two clips detected +// at the same timestamp are still reported in clip order. +#[test] +fn test_same_timestamp_detections_are_reported_in_clip_order() { + let full_audio = Rng::new(11).noise((TIE_FULL_DURATION * SR as f64) as usize, 0.1); + let prefix_audio = &full_audio[..(TIE_PREFIX_DURATION * SR as f64) as usize]; + // Sliding windows 2, 1, 2: the prefix clip is matched after both others. + let clips = vec![ + clip_from_samples(TIE_OTHER_NAME, &Rng::new(12).noise(full_audio.len(), 0.1)), + clip_from_samples(TIE_PREFIX_NAME, prefix_audio), + clip_from_samples(TIE_FULL_NAME, &full_audio), + ]; + let mut audio = silence(6.0, SR); + insert_at(&mut audio, (TIE_START * SR as f64) as usize, &full_audio); + let detector = new_detector(clips, Some(10)).unwrap(); + + let mut events: Vec<(String, f64)> = Vec::new(); + let mut callback = |name: &str, timestamp: f64| events.push((name.to_string(), timestamp)); + detector + .find_clip_in_audio(&mut stream_from_samples("tie", &audio), Some(&mut callback), false) + .unwrap(); + + assert_eq!( + events, + vec![(TIE_PREFIX_NAME.to_string(), TIE_EXPECTED_TIME), (TIE_FULL_NAME.to_string(), TIE_EXPECTED_TIME)] + ); +} + +const THRESHOLD_NAME: &str = "threshold_pattern"; +const THRESHOLD_OTHER_NAME: &str = "threshold_other"; +const THRESHOLD_PATTERN_SAMPLES: usize = 100; +const THRESHOLD_OTHER_SAMPLES: usize = 5000; +/// Amplitude of the quiet copy that puts its correlation peak right at the +/// default `height_min` (0.25) relative to the full-strength copy. +const THRESHOLD_QUIET_SCALE: f32 = 0.2139812; +const THRESHOLD_QUIET_START: usize = 4000; +const THRESHOLD_FULL_START: usize = 12000; +const THRESHOLD_EXPECTED_TIMES: [f64; 2] = [0.499875, 1.499875]; + +/// A clip's FFT size depends only on its own length: an unrelated longer +/// clip with the same sliding window (which alone would need a larger FFT: +/// 16000 + 99 <= 16384 < 16000 + 4999) must not change the f32 rounding of +/// a peak sitting right at the height threshold. +#[test] +fn test_threshold_peak_is_independent_of_other_clips_in_the_group() { + let pattern = Rng::new(0).noise(THRESHOLD_PATTERN_SAMPLES, 0.1); + let other = Rng::new(1000).noise(THRESHOLD_OTHER_SAMPLES, 0.1); + let quiet: Vec = pattern.iter().map(|v| v * THRESHOLD_QUIET_SCALE).collect(); + let mut audio = silence(2.0, SR); + insert_at(&mut audio, THRESHOLD_QUIET_START, &quiet); + insert_at(&mut audio, THRESHOLD_FULL_START, &pattern); + + let run = |clips: Vec| { + let detector = new_detector(clips, Some(2)).unwrap(); + let (peak_times, _) = detector + .find_clip_in_audio(&mut stream_from_samples("threshold", &audio), None, true) + .unwrap(); + detections(&peak_times.unwrap(), THRESHOLD_NAME).to_vec() + }; + + let alone = run(vec![clip_from_samples(THRESHOLD_NAME, &pattern)]); + let with_other = run(vec![ + clip_from_samples(THRESHOLD_NAME, &pattern), + clip_from_samples(THRESHOLD_OTHER_NAME, &other), + ]); + assert_eq!(alone, THRESHOLD_EXPECTED_TIMES); + assert_eq!(with_other, THRESHOLD_EXPECTED_TIMES); +} diff --git a/docs/pattern-matching.md b/docs/pattern-matching.md index 605f386..cdaa48d 100644 --- a/docs/pattern-matching.md +++ b/docs/pattern-matching.md @@ -16,13 +16,13 @@ Raw audio stream (float32 PCM) | Read fixed-size chunks (default 60s) | - Build a per-clip audio section with previous-chunk context + Build one audio section per distinct sliding_window (previous-chunk context) | Loudness-normalize the audio section to -16 dB LUFS | ─── Step 1: Candidate detection (always runs) ─────────────────── | - FFT cross-correlation against each clip + FFT cross-correlation of the shared section against each clip in the group | Peak detection (height >= 0.25, min distance = clip length) | @@ -58,19 +58,19 @@ This produces a `ClipData` struct per clip containing the normalized audio, clip Audio is read as a stream of float32 samples and split into fixed-size chunks (`seconds_per_chunk`, default 60s). Here, `chunk` means the current chunk being processed and `previous_chunk` means the immediately preceding chunk, if one exists. -The raw chunks themselves are not read with overlap. Instead, for each clip, the detector builds an `audio_section`. In the normal case, it prepends the last `sliding_window` seconds from `previous_chunk` to `chunk`, where `sliding_window = ceil(clip_duration_seconds)`. This ensures patterns near chunk boundaries are not missed. +The raw chunks themselves are not read with overlap. Instead, the detector builds one `audio_section` per group of clips with the same `sliding_window`, where `sliding_window = ceil(clip_duration_seconds)`. In the normal case, it prepends the last `sliding_window` seconds from `previous_chunk` to `chunk`. This ensures patterns near chunk boundaries are not missed. The same prepend is applied to every chunk after the first, including a final short chunk. -Each per-clip `audio_section` is loudness-normalized independently to -16 dB LUFS before correlation. +Each `audio_section` is loudness-normalized to -16 dB LUFS before correlation, so a clip's detections depend only on its own section, never on which other clips are loaded. Clips with the same `sliding_window` have identical sections and share one: it is normalized once and transformed once per FFT size, and each clip's cross-correlation then costs a spectrum product with the clip's cached template spectrum and one inverse FFT. The FFT size is the one the section and that clip need on their own (next power of two of `len(audio_section) + len(clip) - 1`), so clips in a group usually share the forward FFT but never change each other's results. ## Step 1: Candidate Detection (FFT Cross-Correlation) This step always runs first for every clip type. Its job is to locate and center potential match positions; it does not decide whether a candidate is a true match — that is Step 2. -For each chunk and each clip: +For each chunk, each sliding-window group's `audio_section` is loaded into a `CorrelationWorkspace` (`load_signal`). Then for each clip in the group: -1. Compute `fft_correlate_full(audio_section, clip)` and take the absolute value. +1. Compute the full cross-correlation with `CorrelationWorkspace::correlate`, using the clip's `CorrelationTemplate` (its time-reversed spectrum, cached per FFT size for the run), and take the absolute value. The result is the same as `fft_correlate_full(audio_section, clip)`; the workspace only avoids repeating the section's forward FFT. 2. Normalize by `max(self_correlation_max, cross_correlation_max)` so the correlation curve is in [0, 1]. 3. Run peak detection with `height >= 0.25` and `distance >= clip_length` (prevents duplicate detections within one clip duration). diff --git a/docs/python.md b/docs/python.md index 31b2d3b..fb01857 100644 --- a/docs/python.md +++ b/docs/python.md @@ -9,7 +9,7 @@ Wheels (Linux x86_64 and arm64, macOS Apple Silicon, Windows x86_64) are attache ```toml # pyproject.toml (uv) [project] -dependencies = ["audio-pattern-detector==0.4.1"] +dependencies = ["audio-pattern-detector==0.4.2"] [[tool.uv.index]] name = "audio-pattern-detector"