Generate noise with a desired ar1 coefficient or with a desired fourier spectrum

rednoise.py

import numpy as np
from scipy.signal import lfilter
import warnings


# helper functions for monte carlo significance tests.
# Aslak Grinsted 2026


def ar1_estimator(x):
    """
    AR1 - Allen and Smith AR(1) model estimation.

    Parameters
    ----------
    x : Time series (univariate).

    Returns
    -------
    g : Lag-one autocorrelation estimate.
    """

    # Ensure x is a 1D array
    x = np.asarray(x).flatten()
    N = len(x)
    m = np.mean(x)
    x = x - m

    # Covariance estimates:
    vr = np.dot(x, x) / N  # zero lag
    c1 = np.dot(x[: N - 1], x[1:N]) / (N - 1)  # lag1

    # Calculate coefficients for quadratic equation
    B = -c1 * N - vr * N**2 - 2 * vr + 2 * c1 - c1 * N**2 + vr * N
    A = vr * N**2
    C = N * (vr + c1 * N - c1)
    D = B**2 - 4 * A * C

    if D > 0:
        g = (-B - np.sqrt(D)) / (2 * A)
    else:
        warnings.warn("Can not place an upperbound on the unbiased AR1.\n" "\t\t -Series too short or too large a trend.", RuntimeWarning)
        g = np.nan

    # Calculate additional outputs if requested: noise innovation variance and mean square
    # mu2 = -1 / N + (2 / N**2) * ((N - g**N) / (1 - g) - g * (1 - g ** (N - 1)) / (1 - g))
    # c0t = vr / (1 - mu2)
    # a = np.sqrt((1 - g**2) * c0t)

    return g


def ar1noise(N, phi, target_sigma=1):
    # Generate AR(1) noise series of length N with lag-1 autocorrelation phi and standard deviation target_sigma
    x = np.random.randn(N) * np.sqrt(1 - phi**2) * target_sigma
    x[0] = x[0] / np.sqrt(1 - phi**2)
    return lfilter([1], [1, -phi], x)


def fftnoise(signal):
    """
    Generate noise with the same power spectrum as the input signal
    by randomizing the phase of its FFT and applying inverse FFT.

    Parameters:
    -----------
    signal : array-like
        Input time series

    Returns:
    --------
    noise : ndarray
        Generated noise with same power spectrum as input
    """
    signal = np.asarray(signal)

    # Compute FFT
    fft_signal = np.fft.fft(signal)

    # Get amplitude (magnitude) and phase
    amplitude = np.abs(fft_signal)

    # Generate random phases uniformly distributed in [0, 2π)
    random_phases = np.random.uniform(0, 2 * np.pi, len(fft_signal))

    # For real-valued output, ensure Hermitian symmetry
    # (only needed if input is real and you want real output)
    if np.isrealobj(signal):
        # DC component (index 0) should have zero phase
        random_phases[0] = 0

        # Nyquist frequency (for even length) should have zero phase
        if len(signal) % 2 == 0:
            random_phases[len(signal) // 2] = 0

        # Mirror the phases for negative frequencies
        random_phases[len(signal) // 2 + 1 :] = -random_phases[1 : len(signal) // 2][::-1]

    # Reconstruct FFT with original amplitude and random phase
    fft_randomized = amplitude * np.exp(1j * random_phases)

    # Apply inverse FFT
    noise = np.fft.ifft(fft_randomized)

    # Return real part if input was real
    if np.isrealobj(signal):
        noise = np.real(noise)

    return noise


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    N = 1000
    phi = 0.9
    target_sigma = 2.0

    y = ar1noise(N, phi, target_sigma)
    yf = fftnoise(y)
    plt.plot(y, label="AR(1) Noise")
    plt.plot(yf, label="FFT Noise")
    plt.legend()
    plt.xlabel("Time")
    plt.ylabel("Value")
    plt.grid()
    plt.show()