Algorithm for colored noise generation

colored_noise.py

"""Generate colored noise."""
from __future__ import absolute_import, division, print_function
import numpy as np
from scipy.fftpack import rfft, irfft
from scipy.signal import fftconvolve
from scipy._lib._util import check_random_state


def _generate_fir_coefficients(exponent, n_samples):
    h = np.empty(n_samples)
    h[0] = 1
    for i in range(1, len(h)):
        h[i] = (0.5 * exponent + i - 1) / i * h[i - 1]
    return h


def generate_noise_with_power_law_spectrum(exponent, shape, fs=1, scale=1,
                                           circular=False, random_state=0):
    """Generate Gaussian noise with power law spectrum.

    This function generates a sample from random process with PSD equal to
    ``scale / (2 * pi * f)**exponent``.

    The output noise is computed as the convolution of white noise with a
    set of specifically generated coefficients as described in [1]_.

    Parameters:
    -----------
    exponent : float
        Exponent of the power spectrum ``PSD(f) ~ (1/f)**exponent``.
    shape : array_like
        The desired shape of the output. The signal changes along the first
        dimension.
    fs : float, optional
        Assumed sampling frequency. Default is 1.
    f_min : float, optional
        Low frequency cutoff. By default is 0, in which case fs / n_samples is
        used.
    scale :  float, optional
        Scale of PSD as defined above. Default is 1.
    circular : bool
        Whether to use a circular convolution. You should analyze whether it
        is going to work well for a particular value of `exponent`.
        For example, it is completely useless when `exponent` is 2
        (results in constant), but might be advantageous for `exponent` equal
        to 1. Default is False.
    random_state : int, None or RandomState
        Seed to use in the numpy random generator. By default is 0.
    
    Returns
    -------
    result : ndarray
        Generated noise with `shape`.

    References
    ----------
    .. [1] N. Jeremy Kasdin, "Discrete Simulation of Colored Noise and
           Stochastic Processes and 1/f^alpha Power Law Noise Generation"
    """
    rng = check_random_state(random_state)
    shape = np.atleast_1d(shape)

    h = _generate_fir_coefficients(exponent, shape[0])
    h_shape = shape.copy()
    h_shape[1:] = 1
    h = h.reshape(h_shape)
    w = rng.randn(*shape)
    scale *= fs ** (1 - exponent)

    if circular:
        noise = irfft(rfft(h, axis=0) * rfft(w, axis=0), axis=0)
    else:
        noise = fftconvolve(h, w, axes=0)[:shape[0]]

    return (scale * fs**(1 - exponent)) ** 0.5 * noise