1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798 |
- 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)
|