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