def fit_log_psd(fql, lc, sim_extra, suff, Dt=None):
"""Calculated log psd for a set of simulated light curves
Args:
fql: array of frequency boundaries
lc: array of simulated light curves of shape (nsim, 3, nt)
sim_extra: dict from simulate_light_curves
suff: e.g. '1' so files are saved as psd__1.*
Dt: if not None, apply aliasing correction
"""
p0 = np.zeros_like(fql[1:]) - 1
fits = []
for tarr, yarr, yerr in lc:
pmod = fqlag.Psd(tarr, yarr, yerr, fql, dt=Dt, log=True)
pfit = fqlag.misc.maximize(pmod, p0)
if np.any(pfit[0] < -10) or np.any(pfit[1] > 10):
p0 = pfit[0]
p0[p0 < -10] = -2
p0[pfit[1] > 10] = -2
pfit = fqlag.misc.maximize(pmod, p0)
fits.append(np.concatenate([pfit[0], pfit[1], [-pfit[2].fun]]))
# shape: nsim, 2*nfq+1 (pars, perr, loglikelihod)
fits = np.array(fits)
psd_model = sim_extra['psd_model']
psd_model[1] = np.log(psd_model[1])
plot_psd(fits, fql, psd_model)
os.system('mkdir -p figures')
plt.savefig('figures/psd__%s.png' % suff)
os.system('mkdir -p npz')
np.savez('npz/psd__%s.npz' % suff, fits=fits, fql=fql, sim_data=sim_extra)
def test_lcxd():
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
tarr, rarr, rerr, sarr, serr, sim = simulate_lc(n, dt, mu, lag)
#fql = np.concatenate([[0.5/n], np.logspace(np.log10(1.5/n), np.log10(0.4), 6)[1:], [1.]])
fql = np.concatenate([[0.1/n], np.logspace(np.log10(2/n), np.log10(0.6), 4)])
p0 = fql[1:]*0 - 1
#p1mod = fqlag.lPsd(tarr, rarr, rerr, fql, )
p1mod = fqlag.Psd(tarr, rarr, rerr, fql, log=True)
#p1 = optimize(p1mod, p0, check_dpar=lambda dp:np.clip(dp, -4, 4))[0]
p1 = maximize(p1mod, p0)
#p1 = [np.array([-1.98476999, -2.10760789, -3.85694557, -6.82487901])]
#p2mod = fqlag.lPsd(tarr[1:], sarr[1:], serr[1:], fql, )
p2mod = fqlag.Psd(tarr[1:], sarr[1:], serr[1:], fql, log=True)
#p2 = optimize(p2mod, p1, check_dpar=lambda dp:np.clip(dp, -4, 4))[0]
p2 = maximize(p2mod, p1[0])
#p2 = [np.array([-2.84762631, -2.16740838, -3.75367259, -6.8628747 ])]
#p0 = np.concatenate([np.min([p1[0],p2[0]], 0)-3, p1[0]*0+0.5])
#cmod = fqlag.lCxd([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, p1[0], p2[0], )
#cmod = fqlag.Cxd([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, p1[0], p2[0], log=True)
p0 = np.concatenate([p1[0], p2[0], np.min([p1[0],p2[0]], 0)-3, p1[0]*0+1])
#cmod = fqlag.lPCxd([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, dt)
cmod = fqlag.PCxd([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, dt, log=True)
check_grad(cmod, p0, 1e-5); return
#print(cmod.loglikelihood(p0));return
print('-----------------\n')
#c = optimize(cmod, p0, maxiter=10, use_bfgs=True, check_dpar=lambda dp:np.clip(dp, -4, 4))
c = maximize(cmod, p0)
def test_cxdRI():
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
tarr, rarr, rerr, sarr, serr, sim = simulate_lc(n, dt, mu, lag)
#fql = np.concatenate([[0.8/n], np.logspace(np.log10(2/n), np.log10(0.4), 6), [0.6]])
fql = np.concatenate([[0.1/n], np.logspace(np.log10(2/n), np.log10(0.6), 4)])
p0 = (fql[1:]**-2)
p1mod = fqlag.Psd(tarr, rarr, rerr, fql)
#p1 = optimize(p1mod, p0)[0]
p1 = maximize(p1mod, p0)
p2mod = fqlag.Psd(tarr, sarr, serr, fql)
#p2 = optimize(p2mod, p1)[0]
p2 = maximize(p2mod, p1)
p0 = np.concatenate([(p1+p2)*0.3, (p1+p2)*0.3])
cmod = fqlag.CxdRI([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, p1, p2)
fqlag.misc.check_grad(cmod, p0);return
#c = optimize(cmod, p0, tol=1e-3, maxiter=20, use_bfgs=True)
c = maximize(cmod, p0)
def test_psi():
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
tarr, rarr, rerr, sarr, serr, sim = simulate_lc(n, dt, mu, lag)
#fql = np.concatenate([[0.5/n], np.logspace(np.log10(2/n), np.log10(0.4), 6), [1.0]])
fql = np.concatenate([[0.1/n], np.logspace(np.log10(2/n), np.log10(0.6), 6)])
p0 = fql[1:]**-2
p1mod = fqlag.Psd(tarr, rarr, rerr, fql)
p1 = maximize(p1mod, p0)[0]
cmod = fqlag.Psi([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, p1)
p0 = np.concatenate([p1*0.1, p1*0])
fqlag.misc.check_grad(cmod, p0)
def test_lppsi():
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
tarr, rarr, rerr, sarr, serr, sim = simulate_lc(n, dt, mu, lag)
#fql = np.concatenate([[0.5/n], np.logspace(np.log10(2/n), np.log10(0.4), 6), [1.0]])
fql = np.concatenate([[0.1/n], np.logspace(np.log10(2/n), np.log10(0.6), 6)])
p0 = np.log(fql[1:]**-2)
#p1mod = fqlag.lPsd(tarr, rarr, rerr, fql)
p1mod = fqlag.Psd(tarr, rarr, rerr, fql, log=True)
p1 = maximize(p1mod, p0)[0]
#cmod = fqlag.lPPsi([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql)
cmod = fqlag.PPsi([tarr, tarr[1:]], [rarr, sarr[1:]], [rerr, serr[1:]], fql, log=True)
p0 = np.concatenate([p1, p1*0-0.5, p1*0])
fqlag.misc.check_grad(cmod, p0, 1e-6); return
c = maximize(cmod, p0)
def test_lpsd():
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
tarr, rarr, rerr, sarr, serr, sim = simulate_lc(n, dt, mu, lag)
fql = np.concatenate([[0.1/n], np.logspace(np.log10(2/n), np.log10(0.4), 4), [1.]])
p0 = np.log(fql[1:]**-2)
#lpmod = fqlag.lPsd(tarr, rarr, rerr, fql, dt)
lpmod = fqlag.Psd(tarr, rarr, rerr, fql, dt, log=True)
#check_grad(lpmod, p0); return
p = maximize(lpmod, p0)
return
y = lpmod.sample(p, 3)
plt.plot(tarr, rarr-rarr.mean())
for x in y:
plt.plot(tarr, x, lw=0.5)
plt.show()
def test_psd():
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
tarr, rarr, rerr, sarr, serr, sim = simulate_lc(n, dt, mu, lag)
fql = np.concatenate([[0.1/n], np.logspace(np.log10(2/n), np.log10(0.4), 4), [1.]])
p0 = (fql[1:]**-1)
pmod = fqlag.Psd(tarr, rarr, rerr, fql)
check_grad(pmod, p0); return
#from IPython import embed; embed();exit(0)
maximize(pmod, p0)
#p = optimize(pmod, p0, maxiter=30, use_bfgs=False)#[0]
#from IPython import embed; embed();exit(0)
return
y = pmod.sample(p, 3)
plt.plot(tarr, rarr-rarr.mean())
for x in y:
plt.plot(tarr, x, lw=0.5)
plt.show()
def test_psd_sim():
# no bias needed. works without the fator of 2
# so not using the factor of 2 is the correct way, and when doing the logs
# no bias correction is needed
nsim = 20
n, dt, mu, lag = 2**8, 1.0, 10000, 1.0
sim = az.SimLC(323)
sim.add_model('broken_powerlaw', [1e-4, -1, -2, 1e-3])
fql = np.concatenate([[0.1/n], np.logspace(np.log10(1.1/n), np.log10(0.4), 7)[1:], [2.]])
p0 = fql[1:]*0+0.1
P = []
for isim in range(1, nsim+1):
sim.simulate(n, dt, mu, 'rms')
tarr, rarr = sim.t[:n], sim.x[:n]
rerr = rarr *0
#rarr = np.random.poisson(rarr)
#rerr = rarr*0 + np.mean(rarr**0.5)
#pmod = fqlag.lPsd(tarr, rarr, rerr, fql)
pmod = fqlag.Psd(tarr, rarr, rerr, fql)
p = maximize(pmod, p0)
P.append(p)
P = np.array(P)
#from IPython import embed; embed();exit(0)
f,fp = sim.normalized_psd
#plt.semilogx(f[1:-1], np.log(fp[1:-1]))
plt.loglog(f[1:-1], fp[1:-1])
fq = (fql[1:] + fql[:-1])/2
plt.errorbar(fq, P.mean(0), P.std(0), fmt='o-')
plt.show()