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_BUCKETS_PER_SEC: f64 = 40.0;
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<Vec<f64>>,
46 span_secs: f64,
47}
48
49fn hop_of(sample_rate: f64) -> usize {
50 ((HOP_SECS * sample_rate).round() as usize).max(1)
51}
52
53fn spectral_frames(planes: &[&[f64]], sample_rate: f64, start_secs: f64) -> Vec<SpectralFrame> {
54 let frame_len = pinned_frame(FRAME_SECS, sample_rate).min(MAX_PINNED_FRAME);
55 let hop = hop_of(sample_rate);
56 let len = planes.iter().map(|p| p.len()).min().unwrap_or(0);
57 (0..len)
58 .step_by(hop)
59 .map(|start| {
60 let end = (start + frame_len).min(len);
61 let mags = planes
62 .iter()
63 .map(|p| {
64 let mut frame = p[start..end].to_vec();
65 frame.resize(frame_len, 0.0);
66 magnitudes(&frame, sample_rate, Some(FRAME_SECS)).0
67 })
68 .collect();
69 SpectralFrame {
70 t_secs: start_secs + start as f64 / sample_rate,
71 mags,
72 span_secs: frame_len as f64 / sample_rate,
73 }
74 })
75 .collect()
76}
77
78pub fn detect(planes: &[&[f64]], sample_rate: f64, start_secs: f64) -> Onsets {
80 let frames = spectral_frames(planes, sample_rate, start_secs);
81 let floor = LEVEL_FRACTION * frames.iter().map(|f| total(&f.mags)).fold(0.0, f64::max);
82 let flux = flux_of(&frames, floor);
83 let onsets = pick_peaks(
84 &flux,
85 &frames,
86 &Energy::of(planes, sample_rate, start_secs),
87 floor,
88 );
89 Onsets {
90 ioi_histogram: ioi_histogram(&onsets, sample_rate),
91 onsets,
92 resolution_secs: hop_of(sample_rate) as f64 / sample_rate,
93 }
94}
95
96fn total(mags: &[Vec<f64>]) -> f64 {
97 mags.iter().flatten().sum()
98}
99
100fn flux_of(frames: &[SpectralFrame], floor: f64) -> Vec<f64> {
103 let opens_on_sound = frames
104 .first()
105 .is_some_and(|f| floor > 0.0 && total(&f.mags) > floor);
106 (0..frames.len())
107 .map(|i| match i {
108 0 => match opens_on_sound {
109 true => total(&frames[0].mags),
110 false => 0.0,
111 },
112 _ => frames[i - 1]
113 .mags
114 .iter()
115 .zip(&frames[i].mags)
116 .map(|(prev, curr)| spectral_flux(prev, curr))
117 .sum(),
118 })
119 .collect()
120}
121
122pub(crate) fn spectral_flux(prev: &[f64], curr: &[f64]) -> f64 {
123 prev.iter().zip(curr).map(|(&p, &c)| (c - p).max(0.0)).sum()
124}
125
126fn pick_peaks(flux: &[f64], frames: &[SpectralFrame], energy: &Energy, floor: f64) -> Vec<Onset> {
129 let mut onsets: Vec<Onset> = Vec::new();
130 let hop_secs = hop_of(energy.sample_rate) as f64 / energy.sample_rate;
131 let window = (ADAPTIVE_WINDOW_SECS / hop_secs).round() as usize;
132 let reach = (MIN_ONSET_GAP_SECS / 2.0 / hop_secs).round().max(1.0) as usize;
133 for i in 0..flux.len() {
134 let lo = i.saturating_sub(window);
135 let hi = (i + window + 1).min(flux.len());
136 let local = &flux[lo..hi];
137 let mean = local.iter().sum::<f64>() / local.len() as f64;
138 let threshold = THRESHOLD_DELTA + THRESHOLD_MULTIPLIER * mean;
139 let near = &flux[i.saturating_sub(reach)..(i + reach + 1).min(flux.len())];
140 if flux[i] <= threshold || flux[i] < floor || near.iter().any(|&v| v > flux[i]) {
141 continue;
142 }
143 let Some(t) = energy.strike(frames[i].t_secs, frames[i].span_secs) else {
144 continue;
145 };
146 if onsets
147 .last()
148 .is_some_and(|o| t - o.t_secs < MIN_ONSET_GAP_SECS)
149 {
150 continue;
151 }
152 onsets.push(Onset {
153 t_secs: t,
154 strength: flux[i],
155 });
156 }
157 onsets
158}
159
160struct Energy {
161 running: Vec<f64>,
162 sample_rate: f64,
163 start_secs: f64,
164 reach: usize,
165 floor: f64,
166}
167
168impl Energy {
169 fn of(planes: &[&[f64]], sample_rate: f64, start_secs: f64) -> Energy {
170 let len = planes.iter().map(|p| p.len()).min().unwrap_or(0);
171 let mut running = Vec::with_capacity(len + 1);
172 running.push(0.0);
173 for n in 0..len {
174 let held = running[running.len() - 1];
175 running.push(held + planes.iter().map(|p| p[n] * p[n]).sum::<f64>());
176 }
177 let reach = (STRIKE_SPAN_SECS * sample_rate).round().max(1.0) as usize;
178 let loudest = (0..len)
179 .map(|n| running[(n + reach).min(len)] - running[n])
180 .fold(0.0, f64::max);
181 Energy {
182 running,
183 sample_rate,
184 start_secs,
185 reach,
186 floor: LEVEL_FRACTION * LEVEL_FRACTION * loudest,
187 }
188 }
189
190 fn over(&self, from: usize, to: usize) -> f64 {
192 let last = self.running.len() - 1;
193 self.running[to.min(last)] - self.running[from.min(last)]
194 }
195
196 fn strike(&self, from_secs: f64, span_secs: f64) -> Option<f64> {
199 let at = |secs: f64| {
200 ((secs - self.start_secs) * self.sample_rate)
201 .round()
202 .max(0.0) as usize
203 };
204 let (from, to) = (at(from_secs), at(from_secs + span_secs));
205 let before = |n: usize| self.over(n.saturating_sub(self.reach), n) + self.floor;
206 let mut best: Option<(f64, usize)> = None;
207 for n in from..to.min(self.running.len() - 1) {
208 let after = self.over(n, n + self.reach);
209 if before(n) == 0.0 {
210 continue;
211 }
212 let ratio = after / before(n);
213 if best.is_none_or(|(r, _)| ratio >= r) {
214 best = Some((ratio, n));
215 }
216 }
217 let (ratio, n) = best?;
218 let carried = before(n) / before(n.saturating_sub(self.reach));
219 (ratio / carried >= JUMP_ENERGY_RATIO)
220 .then(|| self.start_secs + n as f64 / self.sample_rate)
221 }
222}
223
224fn ioi_histogram(onsets: &[Onset], sample_rate: f64) -> Vec<IoiBucket> {
226 let n = (IOI_MAX_SECS * IOI_BUCKETS_PER_SEC).ceil() as usize + 1;
227 let mut counts = vec![0usize; n];
228 for w in onsets.windows(2) {
229 let samples = ((w[1].t_secs - w[0].t_secs) * sample_rate).round();
230 let bucket = (samples * IOI_BUCKETS_PER_SEC / sample_rate).floor() as usize;
231 counts[bucket.min(n - 1)] += 1;
232 }
233 counts
234 .into_iter()
235 .enumerate()
236 .map(|(i, count)| IoiBucket {
237 lo_secs: i as f64 / IOI_BUCKETS_PER_SEC,
238 hi_secs: if i + 1 == n {
239 f64::INFINITY
240 } else {
241 (i + 1) as f64 / IOI_BUCKETS_PER_SEC
242 },
243 count,
244 })
245 .collect()
246}