Matched-filtering

CCF.py

import numpy as np
# from statsmodels.tsa.stattools import ccf
import scipy
from scipy import optimize


def template_fun(x, x_0, W_FWHM, A):
    return A * np.exp(-0.5 * (x - x_0)**2 * (2.35/W_FWHM)**2)


## Assignment ##
x_0 = 423.7
SNR = 5.0
x = np.array(range(1024))
Nr = 100
W_d = 20.7
W_1 = .5 * W_d
W_2 = W_d
W_3 = 1.5 * W_d
amp = 1.


tx = x
T_1 = template_fun(tx, x_0, W_1, amp)
T_2 = template_fun(tx, x_0, W_2, amp)
T_3 = template_fun(tx, x_0, W_3, amp)


def answer(realization, template, verbose=False):
    import matplotlib.pyplot as plt
    cc_loc = []
    for i in range(realization):
        np.random.seed(i)
        noise = np.random.normal(0, 0.5, size=(len(x)))/SNR
        data = amp * np.exp(-0.5 * ((x - x_0)**2) * (2.35/W_d)**2) + noise
        if i == 1 and verbose:
            plt.plot(data)
            plt.show()
        dum = np.correlate(data, template, mode='same')
        _x = np.linspace(-len(dum)/2, len(dum)/2-1, len(dum))
        guess = [0, W_d, 1]     # center of correlation should be around 0 instead of 423.7 because data and template are both centered on 423.7, so their correlation will have ~ 0 lag.
        result = optimize.curve_fit(template_fun, _x, dum, p0=guess)[0]
        if i == 1 and verbose:
            plt.plot(_x, dum)
            plt.plot(_x, template_fun(_x, *result))
            plt.show()
        cc_loc.append(result[0])
    cc_loc = np.array(cc_loc)
    _mean = cc_loc.mean()
    _std = cc_loc.std()
    print "mean: ", _mean
    print "rms  ", np.sqrt(_std**2 + _mean**2)

answer(Nr, T_1)
answer(Nr, T_2)
answer(Nr, T_3)