Skip to main content

sound_to_patch/notes/
mask.rs

1//! Harmonic masks from a note list, and the energy outside them.
2
3use crate::audio::SampleRate;
4use crate::cast;
5use crate::dsp::{Stft, StftParams};
6use crate::units::{Hz, Normalized, Seconds};
7
8use super::NoteList;
9
10/// Settings of the harmonic masks.
11///
12/// The mask of note `n` at frequency `f` and time `t` is
13/// `M_n(f, t) = sum_{h=1..H} exp(-(f - h f_n(t))^2 / (2 sigma^2))`, where
14/// `f_n(t)` follows the note's contour. A cell lies inside a mask when the
15/// value reaches the threshold.
16#[derive(Clone, Copy, Debug, PartialEq)]
17pub struct MaskParams {
18    harmonics: usize,
19    sigma: Option<Hz>,
20    threshold: f64,
21}
22
23impl MaskParams {
24    /// 20 harmonics, `sigma` of one unpadded bin, `fs / frame_len`, and a
25    /// threshold of `exp(-2)`, the mask value two `sigma` from a harmonic.
26    /// That covers the 4-bin main lobe of a Hann window. Zero padding narrows
27    /// the bins but not the lobe, so `sigma` ignores the padding.
28    #[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    /// Sets the number of harmonics `H`. Harmonics at or above the Nyquist
38    /// frequency are skipped.
39    #[must_use]
40    pub fn with_harmonics(self, harmonics: usize) -> Self {
41        Self { harmonics, ..self }
42    }
43
44    /// Sets `sigma`, for example to the unison spread when it is known.
45    #[must_use]
46    pub fn with_sigma(self, sigma: Hz) -> Self {
47        Self {
48            sigma: Some(sigma),
49            ..self
50        }
51    }
52
53    /// Sets the mask value a cell needs to count as inside.
54    #[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/// The masks of one frame: for each bin the note that claims it and
70/// whether a second note claims it too.
71#[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    /// An empty mask for a frame of `bins` bins.
80    #[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    /// Returns the index into the note list of the note with the largest
90    /// mask value at a bin, or `None` outside every mask.
91    #[must_use]
92    pub fn owner(&self, bin: usize) -> Option<usize> {
93        self.owner.get(bin).copied().flatten()
94    }
95
96    /// Returns whether the bin lies inside some mask.
97    #[must_use]
98    pub fn is_masked(&self, bin: usize) -> bool {
99        self.owner(bin).is_some()
100    }
101
102    /// Returns whether two or more notes claim the bin, as octaves and
103    /// fifths do with their shared partials.
104    #[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/// Builds the masks of a note list on the bins of an STFT layout.
117#[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    /// Prepares the masks of `notes` for frames of `stft` at `rate`.
130    #[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    /// Fills `out` with the masks of the notes that sound at `time`, the
151    /// centre of a frame.
152    pub fn frame(&mut self, time: Seconds, out: &mut FrameMask) {
153        out.clear();
154        let bins = out.owner.len().min(self.scratch.len());
155        // Beyond 4 sigma a Gaussian term is below 3.4e-4 of its peak.
156        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/// Returns the share of the signal's power that lies outside every mask.
197///
198/// That share is the bleed of other instruments plus any note the list
199/// missed. Power is `|X|^2` over all STFT cells; a frame where no note
200/// sounds counts as outside. Returns `None` for digital silence.
201#[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    /// The mask at 1 s with bins of 10.77 Hz, 44.1 kHz over 4096 samples.
255    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(&notes(&[45.0]));
267        assert!((0..2049).all(|bin| !single.is_shared(bin)));
268        // 110 Hz sits at bin 10.2 of 10.77 Hz.
269        assert_eq!(single.owner(10), Some(0));
270        let octave = mask_of(&notes(&[45.0, 57.0]));
271        // 220 Hz is harmonic 2 of A2 and harmonic 1 of A3.
272        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        // 110 Hz lies 2.33 Hz above the centre of bin 10.
292        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(&notes(&[45.0, 45.0]));
322        assert_eq!(mask.owner(10), Some(0));
323        assert!(mask.is_shared(10));
324    }
325}