elephant-data/unittest/spectral/multitaper_psd/scripts/generate_test_data.py

import numpy as np
import nitime.algorithms as tsa
import scipy.io


def generate_data(length, coeffs, variance, fs, transients=10):
    """
    Generate model data to test PSD prediction

    Parameters
    ----------
    length : int
        length of generated time series
    coeffs : np.array
        coefficients of auto-regressive model used to generate time series.
        coefficients need to be 1D
    variance : float
        variance of noise used to generate time series from auto-regressive
        model
    fs : float
        sampling frequency of time series data
    transients : int
         number of data points dropped from beginning of time series

    Returns
    -------
    times : np.array
        time axis of generated data
    time_series : np.array
        generated time series data
    freqs : np.array
        frequencies of PSD of generated data
    psd_time_series : np.array
        PSD of generated time series data
    """

    # generate time axis for data to be generated
    times = np.linspace(0, 1 / fs, length)

    # determine order of auto-regressive model to generate data
    order = np.size(coeffs)

    # array to store generated data
    time_series_full = np.zeros(length + transients)

    # generate noise
    noise = np.random.normal(0, variance, length + transients)

    # generate time series data from autoregressive model
    for i in range(length + transients):
        for j in range(order):
            time_series_full[i] += time_series_full[i - j - 1] * coeffs[j]
        time_series_full[i] += noise[i]

    # get rid off transients to finalize generated time series
    time_series = time_series_full[transients:]

    # generate frequencies for PSD
    freqs = np.fft.rfftfreq(length, d=1 / fs)

    # generate PSD of generated time series
    arguments = np.arange(1, order + 1, 1)

    prod_f_arg = np.outer(freqs, arguments)

    exps = np.exp(-2 * np.pi * 1j * prod_f_arg / fs)
    sum_exps = np.matmul(exps, coeffs)

    psd_time_series = variance / (fs * np.abs(1 - sum_exps)**2)

    return times, time_series, freqs, psd_time_series


# Choose parameters, coeffs as in nitime utils.ar_generator default selection
# See: http://nipy.org/nitime/api/generated/nitime.utils.html?highlight=utils%20ar_generator#nitime.utils.ar_generator # noqa
length = 2**10
coeffs = np.array([2.7607, -3.8106, 2.6535, -0.9238])
variance = 1.
fs = 0.1

np.random.seed(1234)
times, time_series, freqs, psd_time_series = generate_data(
    length=length,
    coeffs=coeffs,
    variance=variance,
    fs=fs)

f, psd_mt, nu = tsa.multi_taper_psd(time_series, Fs=fs, NW=4,
                                     jackknife=False, low_bias=False)

assert np.all(freqs == f), "Mismatch of frequencies between ground truth " \
                           "and nitime"

data_folder = '../data/'
np.save(data_folder + 'time_series.npy', time_series)
scipy.io.savemat(data_folder + 'time_series.mat', {"time_series": time_series})
np.save(data_folder + 'psd_nitime.npy', psd_mt)
np.save(data_folder + 'psd_ground_truth.npy', psd_time_series)