Skip to main content

audio_lab/signal/test_signal/
noise.rs

1use ndarray::{Array1, Array2};
2use ndarray_rand::{
3    RandomExt,
4    rand::{SeedableRng, rngs::StdRng},
5    rand_distr::Uniform,
6};
7
8use crate::signal::{SignalError, TimeSignal};
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum Spectrum {
12    White,
13    Pink,
14}
15
16#[derive(Debug, thiserror::Error)]
17pub enum NoiseError {
18    #[error(transparent)]
19    Signal(#[from] SignalError),
20}
21
22#[derive(Debug, Clone, Copy)]
23pub struct NoiseConfig {
24    pub spectrum: Spectrum,
25    pub amplitude: f64,
26    pub sample_rate: f64,
27    pub num_channels: usize,
28    /// Seed for the random number generator. `None` uses a random seed.
29    pub seed: Option<u64>,
30}
31
32impl Default for NoiseConfig {
33    fn default() -> Self {
34        Self {
35            spectrum: Spectrum::White,
36            amplitude: 1.0,
37            sample_rate: 48_000.0,
38            num_channels: 1,
39            seed: None,
40        }
41    }
42}
43
44pub fn generate_noise(
45    num_time_steps: usize,
46    config: NoiseConfig,
47) -> Result<TimeSignal, NoiseError> {
48    let NoiseConfig {
49        spectrum,
50        amplitude,
51        sample_rate,
52        num_channels,
53        seed,
54    } = config;
55
56    let mut rng = match seed {
57        Some(s) => StdRng::seed_from_u64(s),
58        None => StdRng::from_os_rng(),
59    };
60    let data = Array2::random_using(
61        (num_channels, num_time_steps),
62        Uniform::new(-1.0, 1.0).unwrap(),
63        &mut rng,
64    );
65    let mut signal = TimeSignal::new(data, sample_rate)?;
66
67    if spectrum == Spectrum::Pink {
68        let mut freq = signal.into_freq();
69        let filter = Array1::from_shape_fn(freq.num_freq_bins(), |i| ((i + 1) as f64).sqrt());
70        for mut channel in freq.channel_iter_mut() {
71            channel.zip_mut_with(&filter, |bin, scale| *bin /= *scale);
72        }
73        signal = freq.into_time();
74    }
75
76    signal.normalize_peak(amplitude);
77
78    let comment = format!("{spectrum:?} noise signal (amplitude = {amplitude})");
79    signal.set_comment(Some(&comment));
80    Ok(signal)
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn generates_reproducible_white_noise() {
89        let config = NoiseConfig {
90            seed: Some(123),
91            amplitude: 0.5,
92            ..Default::default()
93        };
94        let left = generate_noise(32, config).unwrap();
95        let right = generate_noise(32, config).unwrap();
96        assert_eq!(left, right);
97        let peak = left
98            .iter()
99            .fold(0.0_f64, |peak, sample| peak.max(sample.abs()));
100        assert!((peak - 0.5).abs() < 1e-12);
101    }
102
103    #[test]
104    fn pink_noise_is_supported() {
105        let noise = generate_noise(
106            64,
107            NoiseConfig {
108                spectrum: Spectrum::Pink,
109                seed: Some(123),
110                ..Default::default()
111            },
112        )
113        .unwrap();
114        assert_eq!(noise.num_time_steps(), 64);
115    }
116}