Fluctuation period of a signal with a certain filter width

fluctuation_period.py

import numpy as np
import matplotlib.pyplot as plt

def mean_filt(x, n):
    tent = np.concatenate((np.arange(1, 1+n//2 + n%2), np.arange(1, 1+n//2)[::-1]))
    tent = tent/sum(tent)
    return np.convolve(x, tent, mode='same')


def filt(x, n):
    return mean_filt(x+0.001*np.random.randn(x.size), n)


def get_extrema(x):
    time = np.arange(x.size)[1:-1]
    mini = (x[2:] >= x[1:-1]) & (x[:-2] >  x[1:-1])
    maxi = (x[2:] <  x[1:-1]) & (x[:-2] <= x[1:-1])
    if sum(mini) > sum(maxi):
        first = mini
        second = maxi
    else:
        first = maxi
        second = mini
    return time[first], time[second]


def compute_amplitude(x, i_first, i_second):
    x_first, x_second = x[i_first], x[i_second]
    amplitude = 0.5 * abs(0.5 *(x_first[:-1] + x_first[1:]) - x_second[:-1])
    return amplitude


def get_minmax_data(x, t_first, t_second):
    num_extrema = t_second.size
    i_first  = t_first[:num_extrema]
    i_second = t_second[:num_extrema]
    half_period_1 = (i_second - i_first)[:-1]   # time from extr to extr.
    half_period_2 = (i_first[1:]-i_second[:-1]) # time from extr to extr.
    assert all(half_period_1 > 0.0), half_period_2
    assert all(half_period_2 > 0.0), half_period_2
    period = (half_period_1 + half_period_2)
#     amplitude = compute_amplitude(x, i_first, i_second)
    x_first, x_second = x[i_first], x[i_second]
    amplitude = 0.5 * abs(0.5 *(x_first[:-1] + x_first[1:]) - x_second[:-1])
    t = i_second[:-1]
    return t, amplitude, period


def get_median_period_after_filt(x, threshold=None):
    t_first, t_second = get_extrema(x)
    t, a, p = get_minmax_data(x, t_first, t_second)
    if threshold is None:
        threshold = np.median(a)
    idx = a > threshold
    return np.median(p[idx])


def get_median_period(x, n=10, threshold=None):
    x = filt(x, n)
    return get_median_period_after_filt(x, threshold)


def normalize(x):
    return (x-np.mean(x))/np.std(x)
  
if __name__ == '__main__':
    
    sampling_rate = 1 # Hz
    dt = 1.0/np.float(sampling_rate) # sec.
    periods = [20.0, 70.0, 200.0] # sec.
    weights = [1.0, 1.0, 1.0]
    N = 2 * 60 * 60 # samples
    T = dt * N # sec. (2 hours)
    t = dt * np.arange(N)
    
    signal = np.random.randn(N)
    for p, w in zip(periods, weights):
        signal += w * np.sin(2.0 * (np.pi/p * t + np.random.rand()))

    nf = np.arange(5, 150, 2)
    fluct_periods = [
            get_median_period(signal, n=n)
            for n in nf
    ]

    plt.subplot(211)
    plt.plot(t, signal)
    # plt.xlim(t[0], t[-1])
    plt.xlim(2800, 3500)
    plt.ylabel('Signal')
    plt.xlabel('Time (sec.)')
    plt.subplot(212)
    for p in periods:
        plt.axhline(y=p, color='r', ls='--')
    plt.plot(nf, fluct_periods, label='Estimate')
    plt.plot([], [], 'r--', label='True period')
    plt.legend(loc=0)
    plt.xlim(0, 140)
    plt.ylabel('Period (sec.)')
    plt.xlabel('Filter width (sec.)')
    plt.tight_layout()
    plt.savefig("fluctuation_period.jpg")
    plt.show()