1use crate::audio::SampleRate;
4use crate::cast;
5use crate::dsp::{Stft, StftParams};
6use crate::units::{Hz, Normalized, Seconds};
7
8use super::NoteList;
9
10#[derive(Clone, Copy, Debug, PartialEq)]
17pub struct MaskParams {
18 harmonics: usize,
19 sigma: Option<Hz>,
20 threshold: f64,
21}
22
23impl MaskParams {
24 #[must_use]
29 pub fn new() -> Self {
30 Self {
31 harmonics: 20,
32 sigma: None,
33 threshold: (-2.0_f64).exp(),
34 }
35 }
36
37 #[must_use]
40 pub fn with_harmonics(self, harmonics: usize) -> Self {
41 Self { harmonics, ..self }
42 }
43
44 #[must_use]
46 pub fn with_sigma(self, sigma: Hz) -> Self {
47 Self {
48 sigma: Some(sigma),
49 ..self
50 }
51 }
52
53 #[must_use]
55 pub fn with_threshold(self, threshold: Normalized) -> Self {
56 Self {
57 threshold: threshold.get(),
58 ..self
59 }
60 }
61}
62
63impl Default for MaskParams {
64 fn default() -> Self {
65 Self::new()
66 }
67}
68
69#[derive(Clone, Debug, PartialEq)]
72pub struct FrameMask {
73 owner: Vec<Option<usize>>,
74 best: Vec<f64>,
75 claims: Vec<u8>,
76}
77
78impl FrameMask {
79 #[must_use]
81 pub fn new(bins: usize) -> Self {
82 Self {
83 owner: vec![None; bins],
84 best: vec![0.0; bins],
85 claims: vec![0; bins],
86 }
87 }
88
89 #[must_use]
92 pub fn owner(&self, bin: usize) -> Option<usize> {
93 self.owner.get(bin).copied().flatten()
94 }
95
96 #[must_use]
98 pub fn is_masked(&self, bin: usize) -> bool {
99 self.owner(bin).is_some()
100 }
101
102 #[must_use]
105 pub fn is_shared(&self, bin: usize) -> bool {
106 self.claims.get(bin).is_some_and(|&claims| claims >= 2)
107 }
108
109 fn clear(&mut self) {
110 self.owner.fill(None);
111 self.best.fill(0.0);
112 self.claims.fill(0);
113 }
114}
115
116#[derive(Clone, Debug)]
118pub struct Masker<'a> {
119 notes: &'a NoteList,
120 harmonics: usize,
121 threshold: f64,
122 sigma: f64,
123 bin_hz: f64,
124 nyquist: f64,
125 scratch: Vec<f64>,
126}
127
128impl<'a> Masker<'a> {
129 #[must_use]
131 pub fn new(
132 notes: &'a NoteList,
133 params: MaskParams,
134 stft: StftParams,
135 rate: SampleRate,
136 ) -> Self {
137 let bin_hz = stft.bin_width(rate).get();
138 let lobe_bin = rate.hz().get() / cast::usize_to_f64(stft.frame_len());
139 Self {
140 notes,
141 harmonics: params.harmonics,
142 threshold: params.threshold,
143 sigma: params.sigma.map_or(lobe_bin, Hz::get),
144 bin_hz,
145 nyquist: rate.nyquist().get(),
146 scratch: vec![0.0; stft.bins()],
147 }
148 }
149
150 pub fn frame(&mut self, time: Seconds, out: &mut FrameMask) {
153 out.clear();
154 let bins = out.owner.len().min(self.scratch.len());
155 let reach = 4.0 * self.sigma;
157 for (index, note) in self.notes.notes().iter().enumerate() {
158 if !note.sounds_at(time) {
159 continue;
160 }
161 let f0 = note.f0_at(time).get();
162 let mut touched = (usize::MAX, 0);
163 for h in 1..=self.harmonics {
164 let centre = cast::usize_to_f64(h) * f0;
165 if centre >= self.nyquist {
166 break;
167 }
168 let lo = cast::ceil_to_usize((centre - reach) / self.bin_hz).unwrap_or(0);
169 let hi = cast::floor_to_usize((centre + reach) / self.bin_hz)
170 .unwrap_or(0)
171 .min(bins.saturating_sub(1));
172 for bin in lo..=hi {
173 let distance = cast::usize_to_f64(bin) * self.bin_hz - centre;
174 self.scratch[bin] +=
175 (-distance * distance / (2.0 * self.sigma * self.sigma)).exp();
176 }
177 touched = (touched.0.min(lo), touched.1.max(hi));
178 }
179 if touched.0 > touched.1 {
180 continue;
181 }
182 for bin in touched.0..=touched.1 {
183 let value = std::mem::take(&mut self.scratch[bin]);
184 if value >= self.threshold {
185 out.claims[bin] = out.claims[bin].saturating_add(1);
186 if value > out.best[bin] {
187 out.best[bin] = value;
188 out.owner[bin] = Some(index);
189 }
190 }
191 }
192 }
193 }
194}
195
196#[must_use]
202pub fn out_of_mask_fraction(
203 samples: &[f64],
204 rate: SampleRate,
205 stft: StftParams,
206 notes: &NoteList,
207 params: MaskParams,
208) -> Option<Normalized> {
209 let mut transform = Stft::new(stft);
210 let mut masker = Masker::new(notes, params, stft, rate);
211 let mut mask = FrameMask::new(stft.bins());
212 let (mut outside, mut total) = (0.0, 0.0);
213 for frame in 0..stft.frame_count(samples.len()) {
214 masker.frame(stft.frame_time(frame, rate), &mut mask);
215 for (bin, value) in transform.frame(samples, frame).iter().enumerate() {
216 let power = value.norm_sqr();
217 total += power;
218 if !mask.is_masked(bin) {
219 outside += power;
220 }
221 }
222 }
223 (total > 0.0)
224 .then(|| Normalized::new(outside / total).ok())
225 .flatten()
226}
227
228#[cfg(test)]
229mod tests {
230 use super::*;
231 use crate::notes::Note;
232 use crate::units::MidiPitch;
233
234 fn notes(pitches: &[f64]) -> NoteList {
235 NoteList::new(
236 pitches
237 .iter()
238 .map(|&pitch| {
239 Note::new(
240 Seconds::new(0.0).unwrap(),
241 Seconds::new(10.0).unwrap(),
242 MidiPitch::new(pitch).unwrap(),
243 )
244 .unwrap()
245 })
246 .collect(),
247 )
248 }
249
250 fn mask_of(list: &NoteList) -> FrameMask {
251 mask_with(list, MaskParams::new())
252 }
253
254 fn mask_with(list: &NoteList, params: MaskParams) -> FrameMask {
256 let rate = SampleRate::new(44_100).unwrap();
257 let stft = StftParams::new(4096).unwrap();
258 let mut masker = Masker::new(list, params, stft, rate);
259 let mut mask = FrameMask::new(stft.bins());
260 masker.frame(Seconds::new(1.0).unwrap(), &mut mask);
261 mask
262 }
263
264 #[test]
265 fn an_octave_pair_shares_cells_and_a_single_note_shares_none() {
266 let single = mask_of(¬es(&[45.0]));
267 assert!((0..2049).all(|bin| !single.is_shared(bin)));
268 assert_eq!(single.owner(10), Some(0));
270 let octave = mask_of(¬es(&[45.0, 57.0]));
271 assert!(octave.is_shared(20));
273 assert!(!octave.is_shared(10));
274 }
275
276 #[test]
277 fn a_silent_frame_clears_the_mask_of_the_frame_before() {
278 let list = notes(&[45.0]);
279 let rate = SampleRate::new(44_100).unwrap();
280 let stft = StftParams::new(4096).unwrap();
281 let mut masker = Masker::new(&list, MaskParams::new(), stft, rate);
282 let mut mask = FrameMask::new(stft.bins());
283 masker.frame(Seconds::new(1.0).unwrap(), &mut mask);
284 assert!(mask.is_masked(10));
285 masker.frame(Seconds::new(11.0).unwrap(), &mut mask);
286 assert!((0..stft.bins()).all(|bin| !mask.is_masked(bin)));
287 }
288
289 #[test]
290 fn each_parameter_narrows_the_mask_of_a2() {
291 let a2 = notes(&[45.0]);
293 let default = mask_of(&a2);
294 assert!(default.is_masked(10) && default.is_masked(20));
295 let fundamental = mask_with(&a2, MaskParams::new().with_harmonics(1));
296 assert!(fundamental.is_masked(10) && !fundamental.is_masked(20));
297 let narrow = mask_with(&a2, MaskParams::new().with_sigma(Hz::new(1.0).unwrap()));
298 assert!(!narrow.is_masked(10));
299 let strict = mask_with(
300 &a2,
301 MaskParams::new().with_threshold(Normalized::new(0.99).unwrap()),
302 );
303 assert!(!strict.is_masked(10));
304 }
305
306 #[test]
307 fn a_note_on_one_bin_masks_that_bin() {
308 let bin_hz = 44_100.0 / 4096.0;
309 let pitch = MidiPitch::try_from(Hz::new(10.0 * bin_hz).unwrap()).unwrap();
310 let list = notes(&[pitch.get()]);
311 let params = MaskParams::new()
312 .with_harmonics(1)
313 .with_sigma(Hz::new(0.1).unwrap());
314 let mask = mask_with(&list, params);
315 assert_eq!(mask.owner(10), Some(0));
316 assert!(!mask.is_masked(9) && !mask.is_masked(11));
317 }
318
319 #[test]
320 fn of_two_equal_notes_the_first_owns_the_bin() {
321 let mask = mask_of(¬es(&[45.0, 45.0]));
322 assert_eq!(mask.owner(10), Some(0));
323 assert!(mask.is_shared(10));
324 }
325}