def savePeriodogramPlot(freqs, periods, powers, P, fit_p, subfolder=''):
plt.figure(figsize=(12, 4))
xvals = freqs
plt.xlabel(r'frequency, day$^{-1}$')
if settings.periodogram_units == 'f_log':
plt.xscale('log')
plt.xlabel(r'log frequency, day$^{-1}$')
if settings.periodogram_units == 'p':
xvals = periods
plt.xlabel(r'period, days')
if settings.periodogram_units == 'p_log':
xvals = periods
plt.xscale('log')
plt.xlabel(r'log period, days')
ls_axis = [min(xvals), max(xvals), min(powers), max(powers)]
ls_axis[3] += (ls_axis[3] - ls_axis[2]) / 10.0
plt.axis(ls_axis)
plt.plot(xvals, powers, color="#021c1e", linewidth=1)
plt.plot((P, P), (0, max(powers)),
label='original period, P = ' + str(truncate(P)),
color="#2c7873",
linestyle="--")
if fit_p > 0:
plt.plot((fit_p, fit_p), (0, max(powers)),
label='fitted period, P = ' + str(truncate(fit_p)),
color="#fb6542",
linestyle="-")
plt.ylabel(r'power')
plt.legend(loc='best')
plt.tight_layout()
if len(settings.periodogram_plot) > 0:
ensurePathExists(insertSubfolder(settings.periodogram_plot, subfolder))
plt.savefig(insertSubfolder(settings.periodogram_plot, subfolder))
if settings.periodogram_show == True:
plt.show()
else:
plt.clf()
def runDetector(ts, rvs, ts0, rvs0, errs, Ms, Mps, a, P, e, w, v0, SNR, subfolder='', fit_rotation=False):
G = 8.88677e-10 # AU^3 Mearth^-1 day^-2
if fit_rotation:
out( 'subtracting stellar rotation ...' )
def sine(t, period, amp, phase):
return amp*np.sin(2*np.pi*t/period + phase)
for x in range(1,4): # fit P, P/2, P/3
rotation_period = settings.rotation_period / x
fit_curve = lambda _t, _amp, _phase: sine(_t, settings.rotation_period, _amp, _phase)
par, covm = curve_fit(fit_curve, ts, rvs, p0=[max(rvs), 0])
rv_rot_fit = sine(ts, settings.rotation_period, par[0], par[1])
rvs = rvs - rv_rot_fit
out( 'calculating periodogram ...' )
lsg = LombScarglePyAs(ts, rvs, errs, freqs=np.arange(settings.periodogram_low, settings.periodogram_high, settings.periodogram_step))
freqs, powers, fap = lsg.getPeriodogram()
periods = 1.0 / freqs
# save LSP data
if len(insertSubfolder( settings.periodogram_data , subfolder)) > 0:
savedatafile(insertSubfolder( settings.periodogram_data , subfolder), \
(freqs, periods, powers), \
('freq', 'period', 'power'))
out( 'fitting orbit ...' )
fit_success = True
try:
# get nyquist estimate
mindiff = 1e10 # minimum time difference between points
for i in range(0, len(ts)-1):
val = ts[i+1] - ts[i]
if val < mindiff:
mindiff = val
fitter = RVOrbitFitter(ts, rvs, starmass=Ms, periodrange=[2*mindiff, 7000])#periodrange=[0.5*lsg.bestPeriod, 2.0*lsg.bestPeriod])
fit_Mp, fit_p, fit_e, fit_w, fit_v0 = fitter.getParameters(initial_guess=[Mps, P, e, w, v0])
def afromp(p, m):
return np.power( G*(Ms*332978.9 + m)*p*p / (4*np.pi*np.pi) , 1.0/3.0)
fit_a = (afromp(fit_p[0], fit_Mp[0]), afromp(fit_p[1], fit_Mp[1]))
except Exception as exc:
fit_success = False
print(exc)
tmin = min(ts)
tmax = max(ts)
tmargin = 0 # np.abs(tmax - tmin) / 100
tmin -= tmargin
tmax += tmargin
ts_full = np.arange(tmin, tmax, (tmax-tmin)/settings.rv_curves_points)
if not fit_success:
out( 'WARNING : sorry, Master, I could not fit the orbit T_T' )
# save rv data
if len(insertSubfolder(settings.rv_data, subfolder)) > 0:
savedatafile(insertSubfolder(settings.rv_data, subfolder), \
(ts, rvs, rvs-fitter._curve(ts, Mps, Ms, P, e, w, v0)), \
('time[days]', 'RV[m/s]', 'original_resid'))
if len(insertSubfolder( settings.rv_curves , subfolder)) > 0:
savedatafile(insertSubfolder( settings.rv_curves , subfolder), (ts_full, fitter._curve(ts_full, Mps, Ms, P, e, w, v0)), \
('time[days]', 'original_RV[m/s]'))
rv_full_orig = fitter._curve(ts_full, Mps, Ms, P, e, w, v0)
rv_orig = fitter._curve(ts, Mps, Ms, P, e, w, v0)
rv_full_fit = np.linspace(0, 0, len(ts_full))
rv_fit = np.linspace(0, 0, len(ts))
# params = Mp, p, a, e, w, v0
out( 'producing plots ...' )
saveRVPlot(ts, ts_full, rvs, rv_full_fit, rv_full_orig, rv_fit, rv_orig, Ms, SNR, \
(Mps, P, a, e, w, v0), \
(0,0,0,0,0,0), \
(0,0,0,0,0,0), \
subfolder=subfolder)
# LS periodigram
if len(insertSubfolder( settings.periodogram_plot , subfolder)) > 0:
savePeriodogramPlot(freqs, periods, powers, P, -1, subfolder=subfolder)
return
out('[OUTPUT] Mp = ' + str(truncate(fit_Mp[0])) + u'\t+-\t' + str(truncate(fit_Mp[1])) + '\t(' + str(truncate(100.0*fit_Mp[1]/fit_Mp[0])) + ' %) [earth masses]')
out('[OUTPUT] a = ' + str(truncate(fit_a[0]) ) + u'\t+-\t' + str(truncate(fit_a[1]) ) + '\t(' + str(truncate(100.0*fit_a[1]/fit_a[0]) ) + ' %) [AU]')
out('[OUTPUT] e = ' + str(truncate(fit_e[0]) ) + u'\t+-\t' + str(truncate(fit_e[1]) ) + '\t(' + str(truncate(100.0*fit_e[1]/fit_e[0]) ) + ' %)')
out('[OUTPUT] w = ' + str(truncate(fit_w[0]) ) + u'\t+-\t' + str(truncate(fit_w[1]) ) + '\t(' + str(truncate(100.0*fit_w[1]/fit_w[0]) ) + ' %) [radians]')
out('[OUTPUT] v0 = ' + str(truncate(fit_v0[0])) + u'\t+-\t' + str(truncate(fit_v0[1])) + '\t(' + str(truncate(100.0*fit_v0[1]/fit_v0[0])) + ' %) [m/s]')
out('[OUTPUT] P = ' + str(truncate(fit_p[0]) ) + u'\t+-\t' + str(truncate(fit_p[1]) ) + '\t(' + str(truncate(100.0*fit_p[1]/fit_p[0]) ) + ' %) [days]')
out('[OUTPUT] LS P = ' + str(truncate(lsg.bestPeriod)) + ' [days]\n')
# sigma is assumed to be 1
err_sum = np.sum((fitter._curve(ts, fit_Mp[0], Ms, fit_p[0], fit_e[0], fit_w[0], fit_v0[0]) - rvs)**2)
chi2reduced = err_sum / (len(rvs) - 5)
out( 'chi2 = ' + str(chi2reduced))
residuals = rvs - fitter._curve(ts, Mps, Ms, P, e, w, v0)
out( 'RMS = ' + str(np.sqrt(np.mean(residuals**2))))
# save parameters
if len(insertSubfolder(settings.outputfile , subfolder)) > 0:
ensurePathExists(insertSubfolder(settings.outputfile , subfolder))
with open(insertSubfolder(settings.outputfile , subfolder), 'w+') as f:
f.write('Parameter\tunits\t\toriginal\tdetermined\terrorbar')
f.write('\n\nMplanet\t\t' + '[Mearth]\t' + str(Mps) + '\t\t' + str(truncate(fit_Mp[0])) + '\t\t' + str(truncate(fit_Mp[1])))
f.write('\na\t\t\t' + '[AU]\t\t' + str(a) + '\t\t\t' + str(truncate(fit_a[0])) + '\t\t' + str(truncate(fit_a[1])))
f.write('\ne\t\t\t' + '[-]\t\t\t' + str(e) + '\t\t' + str(truncate(fit_e[0])) + '\t\t' + str(truncate(fit_e[1])))
f.write('\nw\t\t\t' + '[radians]\t' + str(w) + '\t\t' + str(truncate(fit_w[0])) + '\t\t' + str(truncate(fit_w[1])))
f.write('\nv0\t\t\t' + '[m/s]\t\t' + str(v0) + '\t\t\t' + str(truncate(fit_v0[0])) + '\t\t' + str(truncate(fit_v0[1])))
f.write('\nP\t\t\t' + '[days]\t\t' + str(truncate(P)) + '\t\t' + str(truncate(fit_p[0])) + '\t\t' + str(truncate(fit_p[1])))
f.write('\n\nchi2reduced = ' + str(chi2reduced))
f.write('\n\nMstar = ' + str(Ms) + ' [solar masses]')
f.write('\nSNR = ' + str(SNR))
# save rv data
if len(insertSubfolder( settings.rv_data , subfolder)) > 0:
savedatafile(insertSubfolder( settings.rv_data , subfolder), \
(ts, rvs, \
rvs-fitter._curve(ts, Mps, Ms, P, e, w, v0), \
rvs-fitter._curve(ts, fit_Mp[0], Ms, fit_p[0], fit_e[0], fit_w[0], fit_v0[0])), \
('time[days]', 'RV[m/s]', \
'original_resid', \
'fit_resid'))
if len(insertSubfolder( settings.rv_curves , subfolder)) > 0:
savedatafile(insertSubfolder( settings.rv_curves , subfolder), \
(ts_full, fitter._curve(ts_full, Mps, Ms, P, e, w, v0), \
fitter._curve(ts_full, fit_Mp[0], Ms, fit_p[0], fit_e[0], fit_w[0], fit_v0[0]), \
fitter._curve(ts_full, fit_Mp[0], Ms, fit_p[0], fit_e[0], fit_w[0], fit_v0[0])-fitter._curve(ts_full, Mps, Ms, P, e, w, v0)), \
('time[days]', 'original_RV[m/s]', \
'fit_RV[m/s]', \
'difference'))
out( 'producing plots ...' )
# LS periodigram
if len(settings.periodogram_plot) > 0 or settings.periodogram_show:
savePeriodogramPlot(freqs, periods, powers, P, fit_p[0], subfolder=subfolder)
# orbit
if len(settings.rv_plot) > 0 or settings.rv_show:
rv_full_fit = fitter._curve(ts_full, fit_Mp[0], Ms, fit_p[0], fit_e[0], fit_w[0], fit_v0[0])
rv_full_orig = fitter._curve(ts_full, Mps, Ms, P, e, w, v0)
rv_fit = fitter._curve(ts, fit_Mp[0], Ms, fit_p[0], fit_e[0], fit_w[0], fit_v0[0])
rv_orig = fitter._curve(ts, Mps, Ms, P, e, w, v0)
# params = Mp, p, a, e, w, v0
saveRVPlot(ts, ts_full, rvs, rv_full_fit, rv_full_orig, rv_fit, rv_orig, Ms, SNR, \
(Mps, P, a, e, w, v0), \
(fit_Mp[0], fit_p[0], fit_a[0], fit_e[0], fit_w[0], fit_v0[0]), \
(fit_Mp[1], fit_p[1], fit_a[1], fit_e[1], fit_w[1], fit_v0[1]), \
subfolder=subfolder)
def getRVData(self,
progressbar=False,
template='average',
outfunction=None,
times_full=None,
cycle=True):
ts0, rvs0 = self.getRVCurve()
fetcher = S1SpectraFetcher()
time_from = self.settings.time_from
wl_model, flux_model = [], []
lines_wls = []
lines_fluxes = []
if progressbar and outfunction != None:
outfunction('producing profiles ...')
# get lines first
points = []
if progressbar:
from progressbar import ProgressBar, SimpleProgress, ETA
pbar = ProgressBar(widgets=[
' ',
SimpleProgress(), ', ',
ETA()
])
points = pbar(range(0, len(rvs0)))
else:
points = range(0, len(rvs0))
for rvi in points:
# stellar rotation phase
start_rphase = self.settings.start_rphase
rotation_period = self.settings.rotation_period
rphase = start_rphase + (ts0[rvi] - time_from) / rotation_period
# magnetic cycle phase
start_mphase = self.settings.start_mphase
magnetic_period = self.settings.magnetic_period
mphase = start_mphase + (ts0[rvi] - time_from) / magnetic_period
# get profile
wls, fluxes = fetcher.getProfile_old(rphase,
mphase,
rv=rvs0[rvi],
timepoint=ts0[rvi])
#wls, fluxes = fetcher.getProfile(time=ts0[rvi], rphase=rphase, rv=rvs0[rvi])
if self.SNR > 0.0:
# add Gaussian noise with sigma = 1.0/SNR
noise = np.random.normal(0, 1.0 / self.SNR, (len(fluxes)))
fluxes = np.asarray([x + y for x, y in zip(fluxes, noise)])
# save profile to file
if len(self.settings.profiles_path) > 0:
path = self.settings.profiles_path.replace(
r'{N}', str(rvi)).replace(r'{RV}', str(
truncate(rvs0[rvi]))).replace(r'{time}',
str(truncate(ts0[rvi])))
ensurePathExists(path)
lsd_rvs = wls #reader.wlToRv(phase=rphase, wls=wls)
savedatafile(path, \
(wls, lsd_rvs, fluxes), ('wl[A]', 'rv_shift[m/s]', 'flux'), \
comment='RV = ' + str(rvs0[rvi]) + ' [m/s],\tt = ' + str(ts0[rvi]) + ' [days],\tstellar rotation phase: ' + str(rphase))
lines_wls.append(wls)
lines_fluxes.append(fluxes)
self.R = fetcher.R
if progressbar and outfunction != None:
outfunction('computing template ...')
if template == 'average':
# average lines and use it as a model
for rvi in range(0, len(rvs0)):
if rvi == 0:
for wli in range(0, len(lines_wls[rvi])):
wl_model.append(lines_wls[rvi][wli])
flux_model.append(lines_fluxes[rvi][wli])
else:
for wli in range(0, len(lines_wls[rvi])):
flux_model[wli] += lines_fluxes[rvi][wli]
lines_wls[rvi] = np.asarray(lines_wls[rvi])
lines_fluxes[rvi] = np.asarray(lines_fluxes[rvi])
n = len(rvs0)
for i in range(0, len(flux_model)):
flux_model[i] /= n
wl_model = np.asarray(wl_model)
flux_model = np.asarray(flux_model)
elif template == 'model':
wl_model, flux_model = self._getModelSpectrum()
else:
raise Exception(
'Unknown template selected in RVSpectraSimulator.getRVData(). Available options: "model" and "average".'
)
if progressbar and outfunction != None:
outfunction('measuring RVs ...')
points = []
if progressbar:
from progressbar import ProgressBar, SimpleProgress, ETA
pbar = ProgressBar(widgets=[
' ',
SimpleProgress(), ', ',
ETA()
])
points = pbar(rvs0)
else:
points = rvs0
i = 0
rvdata = []
for rv in points:
# measure redshift
val = self.redshiftMeasurer.getRV(lines_wls[i],
lines_fluxes[i],
wl_model,
flux_model,
maxshift=self.maxshift,
R=self.R,
fit_type='linespline',
sfactor=self.sfactor)
rvdata.append(val)
i += 1
return ts0, np.asarray(rvdata), ts0, rvs0
w = settings.w
v0 = settings.v0
SNR = settings.SNR
# read or generate observations time points
times = []
if len(settings.time_source) > 0:
times = np.loadtxt(settings.time_source, usecols=(0, ), unpack=True)
times = np.sort(times)
else:
times = getObsPoints(settings.time_from, settings.time_to,
settings.time_step, settings.time_randomize)
# save time points
ensurePathExists(settings.time_data)
with open(settings.time_data, 'w') as f:
for t in times:
f.write(str(t) + '\n')
tmin = min(times)
tmax = max(times)
ts_full = np.arange(tmin, tmax, (tmax - tmin) / settings.rv_curves_points)
gen = RVSpectraSimulator(Mstar=Ms,
Mplanet=Mps,
a=a,
e=e,
w=w,
v0=v0,
SNR=SNR,
def saveRVPlot(ts,
ts_full,
rvs,
rv_full_fit,
rv_full_orig,
rv_fit,
rv_orig,
Ms,
SNR,
params,
params_fit,
params_err,
subfolder=''):
# params = Mp, p, a, e, w, v0
Mplanet = [params[0], params_fit[0], params_err[0]]
P = [params[1], params_fit[1], params_err[1]]
a = [params[2], params_fit[2], params_err[2]]
e = [params[3], params_fit[3], params_err[3]]
w = [params[4], params_fit[4], params_err[4]]
v0 = [params[5], params_fit[5], params_err[5]]
plt.figure(figsize=(12, 6))
plt.suptitle(r'$M_\bigstar = ' + str(truncate(Ms)) + ' M_\odot$, $SNR=' + str(SNR) + '$\n' + \
'Original orbit: $M_p = ' + str(truncate(Mplanet[0])) + ' M_\oplus$, $P = ' + str(truncate(P[0])) + '$ days, $a = ' + str(truncate(a[0])) + \
'$ AU, $e = ' + str(truncate(e[0])) + '$, $w = '+ str(truncate(w[0])) + '$ rad, $v_0 = '+ str(truncate(v0[0])) + '$ m/s \n' + \
'Best fit result: $M_p = ' + str(truncate(Mplanet[1])) + '\pm' + str(truncate(Mplanet[2])) + ' M_\oplus$, $P = ' \
+ str(truncate(P[1])) + '\pm' + str(truncate(P[2])) + '$ days, $a = ' + str(truncate(a[1])) + '\pm'+ str(truncate(a[2])) + \
'$ AU, $e = ' + str(truncate(e[1])) + '\pm' + str(truncate(e[2])) + '$, $w = '+ str(truncate(w[1])) + '\pm' + str(truncate(w[2])) + \
'$ rad, $v_0 = '+ str(truncate(v0[1])) + '\pm'+ str(truncate(v0[2])) + '$ m/s',\
fontsize=10)
gs = gridspec.GridSpec(3, 1, height_ratios=[4, 1, 1])
gs.update(hspace=0.05)
ax0 = plt.subplot(gs[0])
ax1 = plt.subplot(gs[1], sharex=ax0)
ax2 = plt.subplot(gs[2], sharex=ax0)
ax0.plot(ts_full,
rv_full_orig,
label='original curve',
color="#2c7873",
linestyle="--")
ax0.plot(ts,
rvs,
label='datapoints',
alpha=0.7,
marker="o",
markerfacecolor="#021c1e",
markeredgecolor="none",
linestyle="none")
ax0.plot(ts_full, rv_full_fit, label='fit result', color="#fb6542")
ax0.legend(loc=1)
tmin = min(ts)
tmax = max(ts)
# residuals
ax1.plot((tmin, tmax), (0, 0),
label='original curve',
color="#2c7873",
linestyle="--")
ax1.plot(ts_full, rv_full_fit - rv_full_orig, label='fit', color="#fb6542")
ax1.plot(ts,
rvs - rv_orig,
label='data points',
alpha=0.7,
marker="o",
markerfacecolor="#021c1e",
markeredgecolor="none",
linestyle="none")
ax2.plot((tmin, tmax), (0, 0), label='fit', color="#fb6542")
ax2.plot(ts,
rvs - rv_fit,
label='data points',
alpha=0.7,
marker="o",
markerfacecolor="#021c1e",
markeredgecolor="none",
linestyle="none")
ax0.set_ylabel(r'RV, m/s')
ax2.set_xlabel(r'time, days')
ax0.locator_params(axis='y', nbins=10)
ax1.locator_params(axis='y', nbins=4)
ax2.locator_params(axis='y', nbins=4)
for tick in ax0.get_xticklabels():
tick.set_fontsize(0.0)
for tick in ax1.get_xticklabels():
tick.set_fontsize(0.0)
plt.tight_layout()
if len(settings.rv_plot) > 0:
ensurePathExists(insertSubfolder(settings.rv_plot, subfolder))
plt.savefig(insertSubfolder(settings.rv_plot, subfolder))
if settings.rv_show:
plt.show()
else:
plt.clf()