Imports

import math
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid.inset_locator import inset_axes
import scipy.signal
%matplotlib inline

Plot cosmetics

blue = (0.3, 0.3, 0.8)
green = (0.3, 0.8, 0.3)
pink = (.9, .65, .65)
white = (0.6, 0.6, 0.6)

title_size = 24
axislabel_size = 24
axistick_size = 20
lable_size = 20
legendlabel_size = 16
plotlabel_size = 16

Create artificial time-series

f_sample = 20000
f_nyquist = f_sample/2.
T = 10
n_samples = int(1.*T*f_sample)

# Time
t = 1./f_sample*np.array([i for i in range(n_samples)])

# Frequency
f = np.arange(1./T, f_nyquist+2*1./T, 1./T)

# Signal term
f_0 = round(100*np.random.random_sample(), 2)
f_1 = round(f_0 + (100-f_0)*np.random.random_sample(), 2)
A_0 = 5*(np.random.random_sample()+1)
A_1 = 5*(np.random.random_sample()+1)
bias = 1000*np.random.random_sample()
signal = bias + A_0*np.sin(2*np.pi*f_0*t) + A_1*np.sin(2*np.pi*f_1*t)

# Noise term

# White noise
mu = 15
sigma = 40.91
white_noise = np.random.normal(mu, sigma, n_samples)

# Pink noise
A = n_samples*5
rand_phase = 2*np.pi*np.random.rand(len(f))
pink_noise = np.fft.irfft(A/(f**.5)*np.sin(rand_phase))
                 
noise  = white_noise + pink_noise

# Measured data
data = signal + noise

Plot the data

fig, axes = plt.subplots(2, 2, figsize = (8,8))

def setup(ax):
    plt.sca(ax)
    plt.grid()
    plt.xlabel('Time (s)', size = axislabel_size)
    plt.ylabel('Signal', size = axislabel_size)
    plt.tick_params(labelsize = axistick_size)

# Signal
ax = axes[0,0]
plt.sca(ax)
ax.plot(t, signal, zorder = 100, c = green, lw = 3)
ax.plot(t, bias + A_0*np.sin(2*np.pi*f_0*t), c = (0.2, 0.2, 0.2), lw = 2, zorder = 101, ls = '--', label = str(f_0) + ' Hz')
ax.plot(t, bias + A_1*np.sin(2*np.pi*f_1*t), c = (0.2, 0.2, 0.2), lw = 2, zorder = 101, ls = ':', label = str(f_1) + ' Hz')


plt.legend(loc = 'upper left', fontsize = legendlabel_size)

ax.set_xlim(0, 0.1)
ax.set_ylim(0.9*signal.min(), 1.1*signal.max())

ax.set_title('True signal', size = title_size)
setup(ax)

# White noise
ax = axes[0,1]
ax.plot(t, white_noise, c = white, zorder = 100)
ax.set_xlim(0, 0.1)

ax.set_title('White noise', size = title_size)
setup(ax)

# Pink noise
ax = axes[1,0]
ax.plot(t, pink_noise, c = pink, zorder = 100)
ax.set_xlim(0, 0.1)

ax.set_title('Pink noise', size = title_size)
setup(ax)

# Observed signal
ax = axes[1,1]
ax.plot(t, data, zorder = 100, c = blue)
ax.set_xlim(0, 0.1)

ax.set_title('Observed signal', size = title_size)
setup(ax)
fig.tight_layout()

plt.savefig('signal_time-series.png', dpi = 300)

plt.show()

Full-data periodogram

from scipy.signal import periodogram

fig = plt.figure(figsize = (8,6))
f, P_dens = periodogram(data, f_sample)

plt.title('Periodogram of data', size = title_size)
plt.loglog(f, P_dens, zorder = 10, c = (0.2, 0.2, 0.2))
plt.grid()
#plt.plot([f_0, f_0], [0.0000001, 10000], ls = '--', c = 'black', zorder = 0)
#plt.plot([f_1, f_1], [0.0000001, 10000], ls = '--', c = 'black', zorder = 0)


plt.xlabel('Frequency (Hz)', size = axislabel_size)
plt.ylabel('Power (signal$^{2}$)', size = axislabel_size)
plt.ylim(.000001, 2000)

plt.savefig('periodogram.png', dpi = 300)
plt.show()

Welch's method

window_width = 2
window_overlap = 0.25
windows = []
num_windows = int(T/(window_width*(1-window_overlap)))-1
for i in range(num_windows):
    windows.append((int((1-window_overlap)*f_sample*window_width*i), (int((1-window_overlap)*window_width*f_sample*i + window_width*f_sample))))

def get_window(data, window):
    i = window[0]
    j = window[1]
    return data[i:j]

def window_data(data, window):
    data_sub = data[window[0]:window[1]]
    
    N = window[1] - window[0]
    window_function = 0.5*(1-np.cos(2*np.pi/(N-1)*np.arange(0, N, 1)))
    
    return data_sub*window_function

Plot data, windowed data, window function

window = windows[0]
N = window[1] - window[0]

# Begin plot

fig = plt.figure(figsize = (8, 6))

plt.plot(t[window[0]:window[1]], get_window(data, window), color = (0.2, 0.2, 0.2), label = 'Data', zorder = 101)
plt.plot(t[window[0]:window[1]], np.mean(get_window(data, window))*0.5*(1-np.cos(2*np.pi/(N-1)*np.arange(0, N, 1))),\
        ls = '--', lw = 3, c = (0.6, 0.6, 0.6), zorder = 102, label = 'Hanning window')

plt.plot(t[window[0]:window[1]], window_data(data, window), color = (0.4, 0.4, 0.4), label = 'Windowed data', zorder = 101)


plt.grid()
plt.tick_params(labelsize = axistick_size)
leg = plt.legend(fontsize = legendlabel_size)
leg.set_zorder(200)
for legobj in leg.legendHandles:
    legobj.set_linewidth(4.0)
    
plt.xlabel('Time (s)', size = axislabel_size)
plt.ylabel('Signal', size = axislabel_size)

plt.savefig('window_function.png', dpi = 300)

plt.show()

for i in range(num_windows):
    windowed_data = window_data(data, windows[i])
    f_temp, P_dens_temp = periodogram(windowed_data, f_sample)
    if i == 0:
        P_dens_total = P_dens_temp
    else:
        P_dens_total += P_dens_temp

P_dens_total /= num_windows

fig, ax = plt.subplots(1, figsize = (8,6))

plt.title('Periodogram of data', size = title_size)
plt.loglog(f, P_dens, zorder = 9, c = (0.2, 0.2, 0.2))
plt.loglog(f_temp, P_dens_total, zorder = 10, c = (0.6, 0.6, 0.6))
#plt.plot([f_0, f_0], [0.0000001, 1000000], ls = '--', c = 'black', zorder = 0)
#plt.plot([f_1, f_1], [0.0000001, 1000000], ls = '--', c = 'black', zorder = 0)

plt.ylim(.000001, 1000000)

plt.tick_params(labelsize = axistick_size)

plt.grid()

plt.xlabel('Frequency (Hz)', size = axislabel_size)
plt.ylabel('Power (signal$^{2}$)', size = axislabel_size)

plt.savefig('welch_periodogram.png', dpi = 300)

plt.show()

Remove random data points

import random

n_remove = int(n_samples*0.75)
keep_list = np.sort(np.array(random.sample(range(0, n_samples), n_samples - n_remove)))

t_sub = t[keep_list].flatten()
signal_sub = signal[keep_list.flatten()]
data_sub = data[keep_list].flatten()

fig, ax = plt.subplots(211)

plt.title('Down sampled signal')
plt.plot(t_sub, signal_sub, marker = 'o', lw = 1)
plt.plot(t, signal, zorder = 1, lw = 5, c = (0.8,0.2,0.2))
plt.xlim(0, 0.05)
plt.show()

plt.title('Down sampled measurements')
plt.plot(t_sub, data_sub, lw = 1, marker = 'o')
plt.plot(t, data, zorder = 1, lw = 5, c = (0.8,0.2,0.2))
plt.xlim(0,.01)
plt.show()

Lomb-Scargle periodogram

from astropy.stats import LombScargle

f, P = LombScargle(t_sub, data_sub).autopower()

plt.loglog(f, P, zorder = 100)
plt.plot([f_0, f_0], [.0000000000001, 100], ls = '--')
plt.show()