using celerite to identify gran+osc. components

[skgrunblatt]

import celerite
from celerite import terms
from celerite.modeling import Model, ConstantModel

#set the GP parameters

Q = 1.0 / np.sqrt(2.0)
w0 = muhz2omega(13)
S0 = np.var(y) / (w0*Q)
kernel = terms.SHOTerm(log_S0=np.log(S0), log_Q=np.log(Q), log_omega0=np.log(w0),
bounds=[(-25, 0), (-15, 15), (np.log(muhz2omega(5)), np.log(muhz2omega(50)))])
kernel.freeze_parameter("log_Q") #to make it aperiodic

Q = 1.0 / np.sqrt(2.0)
w0 = muhz2omega(81.0)
S0 = np.var(y) / (w0*Q)
kernel += terms.SHOTerm(log_S0=np.log(S0), log_Q=np.log(Q), log_omega0=np.log(w0),
bounds=[(-25, 0), (-15, 15), (np.log(muhz2omega(20)), np.log(muhz2omega(1000)))])
kernel.freeze_parameter("terms[1]:log_Q") #to make it aperiodic

Q = np.exp(3.0)
w0 = muhz2omega(200) #peak of oscillations at ~220 muhz
S0 = np.var(y) / (w0*Q)
kernel += terms.SHOTerm(log_S0=np.log(S0), log_Q=np.log(Q), log_omega0=np.log(w0),
bounds=[(-40, 0), (0.5, 4.2), (np.log(muhz2omega(60)), np.log(muhz2omega(800)))])

lwn = np.log(np.mean(yerr**2)/len(t))

kernel += terms.JitterTerm(log_sigma=-10, bounds=[(-20,20)])

gp = celerite.GP(kernel)
gp.compute(t, yerr)
print("Initial log likelihood: {0}".format(gp.log_likelihood(y)))
gp.get_parameter_dict()

#find max likelihood params

from scipy.optimize import minimize

def neg_log_like(params, y, gp):
gp.set_parameter_vector(params)
ll = gp.log_likelihood(y)
if not np.isfinite(ll):
return 1e10
return -ll

initial_params = gp.get_parameter_vector()
bounds = gp.get_parameter_bounds()

r = minimize(neg_log_like, initial_params, method="L-BFGS-B", bounds=bounds, args=(y, gp))
gp.set_parameter_vector(r.x)
print(r)

#Plot GP fit PSD on top of PSD of data
psd = gp.kernel.get_psd(omega)

#plot individual components
plt.plot(freq, power_ls, lw=0.5)
for k in gp.kernel.terms:
print(k)
plt.plot(freqmuhz, 1e12*k.get_psd(omega)/(2.*np.pi), "--", color='orange')

white_noise = (yerr**2 + gp.kernel.jitter) / len(t)
plt.axhline(1e12white_noise / (2np.pi), ls='--', color='orange')

#plot combined model
plt.plot(freqmuhz, 1e12*(gp.kernel.get_psd(omega) + white_noise) / (2*np.pi), color='r')

plt.yscale("log")
plt.xscale('log')
plt.xlim(2,5000)
plt.xlabel("muhz")
plt.ylabel("power (ppm^2/muhz)");

psd_interp = np.interp(freq, freqmuhz, (1e12*gp.kernel.terms[0].get_psd(omega)/(2.np.pi) + 1e12gp.kernel.terms[1].get_psd(omega)/(2.*np.pi)))

plt.plot(freq,power_ls)
plt.plot(freq,psd_interp)
plt.yscale("log")
plt.xscale('log')
plt.xlim(2,2000)
#plt.ylim(1e-50, 1)
plt.xlabel("muhz")
plt.ylabel("power (ppm^2/muhz)");

bkg_corr = power_ls - psd_interp

plt.clf()
plt.plot(freq,bkg_corr)
plt.yscale("log")
plt.xscale('log')
plt.xlim(150,200)
plt.ylim(1, 3000)
plt.xlabel("muhz")
plt.ylabel("power (ppm^2/muhz)");

#identify expected dnu and plot

acf = np.correlate(power_ls, power_ls, 'same')

acf_corr = np.correlate(bkg_corr, bkg_corr, 'same')
freq_acf = np.linspace(-freq[-1],freq[-1],len(freq))

dnu_exp = 0.267 * omega2muhz((np.e**(gp.kernel.terms[2].log_omega0))) ** 0.764

plt.plot(freq_acf, acf)
plt.plot(freq_acf, acf_corr)
plt.axvline(dnu_exp, ls='--', color='k')
plt.axvspan(0.85dnu_exp, 1.15dnu_exp, color='gray', zorder=0)
#plt.yscale("log")
#plt.xscale('log')
plt.xlim(0,20)
plt.ylim(0, 4e7)
plt.xlabel("muhz")
#plt.ylabel("power (ppm^2/muhz)");

print(gp.kernel.terms[2].log_omega0)
print('numax: ', omega2muhz((np.e**(gp.kernel.terms[2].log_omega0))),'dnu_exp: ',dnu_exp)