Generate noise in Python with a specific colour / power spectral density

noise_generator.py

# 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

Great content and great blog post! Very clean and clear.

Regarding the integration of noise to 0, it definely is due, like you guessed, to the first coefficient of the DFT, set to 0. I guess it is always easier to assume noise as 0-meaned The higher the variance, the wider the curves would be (highest would be higher, lowest would be lower), which mean you deviation from the true value (the mean), would be bigger - but on average you are always around the mean.

I'd be curious to know what happens if you shift (roll actually) your samples, say by half the length of the signal...! (in terms of integral and spectrum)

Regarding the noise distribution turning gaussian after fourier-shaping, I'd go with the central limit theorem : after all, a fourier coefficients is a sum of random samples in this case, and changing the power to each frequency just adds some randomness. I'd add that it is actually a miracle that the FT and IFT journey gives back the original distribution. Maybe try shaping the phase ? Or use few samples to step away from CLT and average several tries (ie, monte-carlo) to see if the original distribution is preserved ?