1use crate::measure::spectrum::{MAX_PINNED_FRAME, magnitudes, pinned_frame};
4
5const FRAME_SECS: f64 = 0.023;
7const HOP_SECS: f64 = FRAME_SECS / 2.0;
8const MIN_ONSET_GAP_SECS: f64 = 0.05;
10const ADAPTIVE_WINDOW_FRAMES: usize = 10;
12const THRESHOLD_MULTIPLIER: f64 = 1.5;
13const THRESHOLD_DELTA: f64 = 1e-6;
14const IOI_BUCKET_SECS: f64 = 0.025;
15const IOI_MAX_SECS: f64 = 2.0;
16const LEVEL_FRACTION: f64 = 0.03;
18const STRIKE_SPAN_SECS: f64 = 0.012;
19const JUMP_ENERGY_RATIO: f64 = 4.0;
21
22#[derive(Clone, Copy, Debug, PartialEq)]
23pub struct Onset {
24 pub t_secs: f64,
25 pub strength: f64,
26}
27
28#[derive(Clone, Copy, Debug, PartialEq)]
29pub struct IoiBucket {
30 pub lo_secs: f64,
31 pub hi_secs: f64,
32 pub count: usize,
33}
34
35#[derive(Clone, Debug, PartialEq)]
36pub struct Onsets {
37 pub onsets: Vec<Onset>,
38 pub ioi_histogram: Vec<IoiBucket>,
39 pub resolution_secs: f64,
40}
41
42struct SpectralFrame {
43 t_secs: f64,
44 mags: Vec<f64>,
45 span_secs: f64,
46}
47
48fn spectral_frames(samples: &[f64], sample_rate: f64, start_secs: f64) -> Vec<SpectralFrame> {
49 let frame_len = pinned_frame(FRAME_SECS, sample_rate).min(MAX_PINNED_FRAME);
50 let hop = ((HOP_SECS * sample_rate).round() as usize).max(1);
51 (0..samples.len())
52 .step_by(hop)
53 .map(|start| {
54 let end = (start + frame_len).min(samples.len());
55 let (mags, ..) = magnitudes(&samples[start..end], sample_rate, Some(FRAME_SECS));
56 SpectralFrame {
57 t_secs: start_secs + start as f64 / sample_rate,
58 mags,
59 span_secs: frame_len as f64 / sample_rate,
60 }
61 })
62 .collect()
63}
64
65pub fn detect(samples: &[f64], sample_rate: f64, start_secs: f64) -> Onsets {
66 let frames = spectral_frames(samples, sample_rate, start_secs);
67 let floor = LEVEL_FRACTION * frames.iter().map(|f| total(&f.mags)).fold(0.0, f64::max);
68 let flux = flux_of(&frames, floor);
69 let onsets = pick_peaks(
70 &flux,
71 &frames,
72 &Energy::of(samples, sample_rate, start_secs),
73 floor,
74 );
75 Onsets {
76 ioi_histogram: ioi_histogram(&onsets),
77 onsets,
78 resolution_secs: HOP_SECS,
79 }
80}
81
82fn total(mags: &[f64]) -> f64 {
83 mags.iter().sum()
84}
85
86fn flux_of(frames: &[SpectralFrame], floor: f64) -> Vec<f64> {
89 let opens_on_sound = frames
90 .first()
91 .is_some_and(|f| floor > 0.0 && total(&f.mags) > floor);
92 (0..frames.len())
93 .map(|i| match i {
94 0 => match opens_on_sound {
95 true => total(&frames[0].mags),
96 false => 0.0,
97 },
98 _ => spectral_flux(&frames[i - 1].mags, &frames[i].mags),
99 })
100 .collect()
101}
102
103fn spectral_flux(prev: &[f64], curr: &[f64]) -> f64 {
104 prev.iter().zip(curr).map(|(&p, &c)| (c - p).max(0.0)).sum()
105}
106
107fn pick_peaks(flux: &[f64], frames: &[SpectralFrame], energy: &Energy, floor: f64) -> Vec<Onset> {
110 let mut onsets: Vec<Onset> = Vec::new();
111 for i in 0..flux.len() {
112 let lo = i.saturating_sub(ADAPTIVE_WINDOW_FRAMES);
113 let hi = (i + ADAPTIVE_WINDOW_FRAMES + 1).min(flux.len());
114 let local = &flux[lo..hi];
115 let mean = local.iter().sum::<f64>() / local.len() as f64;
116 let threshold = THRESHOLD_DELTA + THRESHOLD_MULTIPLIER * mean;
117 let reach = (MIN_ONSET_GAP_SECS / HOP_SECS).round().max(1.0) as usize;
118 let near = &flux[i.saturating_sub(reach)..(i + reach + 1).min(flux.len())];
119 if flux[i] <= threshold || flux[i] < floor || near.iter().any(|&v| v > flux[i]) {
120 continue;
121 }
122 let Some(t) = energy.strike(frames[i].t_secs, frames[i].span_secs) else {
123 continue;
124 };
125 if onsets
126 .last()
127 .is_some_and(|o| t - o.t_secs < MIN_ONSET_GAP_SECS)
128 {
129 continue;
130 }
131 onsets.push(Onset {
132 t_secs: t,
133 strength: flux[i],
134 });
135 }
136 onsets
137}
138
139struct Energy {
140 running: Vec<f64>,
141 sample_rate: f64,
142 start_secs: f64,
143 reach: usize,
144 floor: f64,
145}
146
147impl Energy {
148 fn of(samples: &[f64], sample_rate: f64, start_secs: f64) -> Energy {
149 let mut running = Vec::with_capacity(samples.len() + 1);
150 running.push(0.0);
151 for s in samples {
152 let held = running[running.len() - 1];
153 running.push(held + s * s);
154 }
155 let reach = (STRIKE_SPAN_SECS * sample_rate).round().max(1.0) as usize;
156 let loudest = (0..samples.len())
157 .map(|n| running[(n + reach).min(samples.len())] - running[n])
158 .fold(0.0, f64::max);
159 Energy {
160 running,
161 sample_rate,
162 start_secs,
163 reach,
164 floor: LEVEL_FRACTION * LEVEL_FRACTION * loudest,
165 }
166 }
167
168 fn over(&self, from: usize, to: usize) -> f64 {
170 let last = self.running.len() - 1;
171 self.running[to.min(last)] - self.running[from.min(last)]
172 }
173
174 fn strike(&self, from_secs: f64, span_secs: f64) -> Option<f64> {
177 let at = |secs: f64| {
178 ((secs - self.start_secs) * self.sample_rate)
179 .round()
180 .max(0.0) as usize
181 };
182 let (from, to) = (at(from_secs), at(from_secs + span_secs));
183 let mut best: Option<(f64, usize)> = None;
184 for n in from..to.min(self.running.len() - 1) {
185 let before = self.over(n.saturating_sub(self.reach), n) + self.floor;
186 let after = self.over(n, n + self.reach);
187 if before == 0.0 {
188 continue;
189 }
190 let ratio = after / before;
191 if best.is_none_or(|(r, _)| ratio >= r) {
192 best = Some((ratio, n));
193 }
194 }
195 let (ratio, n) = best?;
196 (ratio >= JUMP_ENERGY_RATIO).then(|| self.start_secs + n as f64 / self.sample_rate)
197 }
198}
199
200fn ioi_histogram(onsets: &[Onset]) -> Vec<IoiBucket> {
201 let n = (IOI_MAX_SECS / IOI_BUCKET_SECS).ceil() as usize + 1;
202 let mut counts = vec![0usize; n];
203 for w in onsets.windows(2) {
204 let ioi = w[1].t_secs - w[0].t_secs;
205 counts[((ioi / IOI_BUCKET_SECS) as usize).min(n - 1)] += 1;
206 }
207 counts
208 .into_iter()
209 .enumerate()
210 .map(|(i, count)| IoiBucket {
211 lo_secs: i as f64 * IOI_BUCKET_SECS,
212 hi_secs: if i + 1 == n {
213 f64::INFINITY
214 } else {
215 (i + 1) as f64 * IOI_BUCKET_SECS
216 },
217 count,
218 })
219 .collect()
220}