-
-
Save raytroop/8818d2238dcfd0706e57b82b5bf7fec2 to your computer and use it in GitHub Desktop.
Generate noise in Python with a specific colour / power spectral density
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Copyright 2022 Ioces Pty. Ltd. | |
| # | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy | |
| # of this software and associated documentation files (the "Software"), to deal | |
| # in the Software without restriction, including without limitation the rights | |
| # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | |
| # copies of the Software, and to permit persons to whom the Software is | |
| # furnished to do so, subject to the following conditions: | |
| # | |
| # The above copyright notice and this permission notice shall be included in | |
| # all copies or substantial portions of the Software. | |
| # | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE. | |
| import numpy as np | |
| class NoiseGenerator: | |
| """Generates gaussian noise with a specified power spectral density. | |
| Noise is generated by first creating a gaussian white base noise in the | |
| frequency domain, and then shaping it using a colouring function. Various | |
| colouring functions are provided for convenience. | |
| Examples | |
| -------- | |
| >>> ng = NoiseGenerator() | |
| Generate 100000 white noise points sampled at 1 kHz with a PSD of 0.1. | |
| >>> white = ng.generate(1e-3, 100000, colour=ng.white(0.1)) | |
| Generate some other colours of noise at a sampling rate of 1 MHz. | |
| >>> pink = ng.generate(1e-6, 1000000, colour=ng.pink(10.0)) | |
| >>> blue = ng.generate(1e-6, 1000000, colour=ng.blue()) | |
| Use a custom piecewise colouring function to generate specific noise. | |
| >>> frequencies = [90.0, 100.0, 110.0, 450.0, 500.0, 550.0] | |
| >>> psds = [0.1, 10.0, 0.01, 0.01, 2.0, 0.001] | |
| >>> colour = ng.piecewise_logarithmic(frequencies, psds) | |
| >>> custom = ng.generate(1e-4, 1000000, colour=colour) | |
| """ | |
| rng = np.random.default_rng() | |
| def generate(self, dt, n, colour=None): | |
| """Generates uniformly sampled noise of a particular colour. | |
| Parameters | |
| ---------- | |
| dt : float | |
| Sampling period, in seconds. | |
| n : int | |
| Number of samples to generate. | |
| colour : function, optional | |
| Colouring function that specifies the PSD at a given frequency. If | |
| not specified, the noise returned will be white Gaussian noise with | |
| a PSD of 1.0 across the entire frequency range. | |
| Returns | |
| ------- | |
| numpy.array | |
| Length `n` array of sampled noise. | |
| """ | |
| f, x_f = self._base_noise(dt, n) | |
| if colour: | |
| x_f *= np.sqrt(colour(f)) | |
| return np.fft.irfft(x_f) | |
| @staticmethod | |
| def white(scale=1.0): | |
| """Creates a white noise colouring function. | |
| Parameters | |
| ---------- | |
| scale : float, optional | |
| Multiplier to adjust the scale of the white noise. | |
| Returns | |
| ------- | |
| function | |
| White noise colouring function that can be used with `generate()`. | |
| The function returned will be :math:`y(f) = s`, where :math:`s` is | |
| `scale`. | |
| """ | |
| return lambda f: scale | |
| @staticmethod | |
| def pink(scale=1.0): | |
| """Creates a pink noise colouring function. | |
| Parameters | |
| ---------- | |
| scale : float, optional | |
| Multiplier to adjust the scale of the pink noise. | |
| Returns | |
| ------- | |
| function | |
| Pink noise colouring function that can be used with `generate()`. | |
| The function returned will be :math:`y(f) = s / f`, where :math:`s` | |
| is `scale`. At f = 0, the function will simply return 0.0. | |
| """ | |
| return lambda f: scale / np.where(f == 0.0, np.inf, f) | |
| @staticmethod | |
| def brownian(scale=1.0): | |
| """Creates a brownian noise colouring function. | |
| Parameters | |
| ---------- | |
| scale : float, optional | |
| Multiplier to adjust the scale of the brownian noise. | |
| Returns | |
| ------- | |
| function | |
| Brownian noise colouring function that can be used with | |
| `generate()`. The function returned will be | |
| :math:`y(f) = s / f ^ 2`, where :math:`s` is `scale`. At f = 0, the | |
| function will simply return 0.0. | |
| """ | |
| return lambda f: scale / np.where(f == 0.0, np.inf, f**2) | |
| @staticmethod | |
| def blue(scale=1.0): | |
| """Creates a blue noise colouring function. | |
| Parameters | |
| ---------- | |
| scale : float, optional | |
| Multiplier to adjust the scale of the blue noise. | |
| Returns | |
| ------- | |
| function | |
| Blue noise colouring function that can be used with `generate()`. | |
| The function returned will be :math:`y(f) = s * f`, where :math:`s` | |
| is `scale`. | |
| """ | |
| return lambda f: scale * f | |
| @staticmethod | |
| def violet(scale=1.0): | |
| """Creates a blue noise colouring function. | |
| Parameters | |
| ---------- | |
| scale : float, optional | |
| Multiplier to adjust the scale of the violet noise. | |
| Returns | |
| ------- | |
| function | |
| Violet noise colouring function that can be used with `generate()`. | |
| The function returned will be :math:`y(f) = s * f ^ 2`, where | |
| :math:`s` is `scale`. | |
| """ | |
| return lambda f: scale * f**2 | |
| # Some common aliases for the various colours of noise | |
| brown = brownian | |
| red = brownian | |
| azure = blue | |
| purple = violet | |
| @staticmethod | |
| def piecewise_logarithmic(frequencies, psds): | |
| """Creates a custom piecewise colouring function | |
| Parameters | |
| ---------- | |
| frequencies : numpy.array | |
| Array of frequencies, in Hz | |
| psds : numpy.array | |
| Array of PSDs | |
| Returns | |
| ------- | |
| function | |
| Custom noise colouring function that can be used with `generate()`. | |
| The function is linearly interpolated in log space for the given | |
| frequencies and PSDs. Values outside the range of frequencies given | |
| are set to the PSD of the closest endpoint. | |
| """ | |
| # Convert to log space | |
| log_frequencies = np.log(frequencies) | |
| log_psds = np.log(psds) | |
| # Create a closure for our colour function that suppresses the warning | |
| # about np.log(0) (which correctly returns -np.inf anyway) | |
| def colour(f): | |
| with np.errstate(divide="ignore"): | |
| return np.exp(np.interp(np.log(f), log_frequencies, log_psds)) | |
| return colour | |
| def _base_noise(self, dt, n): | |
| """Produces a frequency domain representation of uniformly sampled | |
| Gaussian white noise. | |
| """ | |
| # Calculate random frequency components | |
| f = np.fft.rfftfreq(n, dt) | |
| x_f = self.rng.normal(0, | |
| 0.5, len(f)) + 1j * self.rng.normal(0, 0.5, len(f)) | |
| x_f *= np.sqrt(n / dt) | |
| # Ensure our 0 Hz and Nyquist components are purely real | |
| x_f[0] = np.abs(x_f[0]) | |
| if len(f) % 2 == 0: | |
| x_f[-1] = np.abs(x_f[-1]) | |
| return f, x_f |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment