diff --git a/examples/nbody/README.md b/examples/nbody/README.md index d70a6f70..a0c0620c 100644 --- a/examples/nbody/README.md +++ b/examples/nbody/README.md @@ -1,7 +1,15 @@ +# N-body Retrievals + Coming soon to EXOTIC... -In the meantime checkout: https://github.com/pearsonkyle/Nbody-ai +## Prior estimate + +The lomb-scargle periodogram from our ephemeris fitting code can be used to estimate a prior for an N-body retrieval. Transit timing variations (TTV) alone are not sufficient to uniquely constrain a single solution, instead multiple modes can exist which give the same amplitude/periodicity of perturbation. The distinguishing factor will rely on radial velocity measurements but none the less this approach can help limit the search space for potential planets. Below is a figure highlighting results from a grid of N-body simulations where we construct a mask based on the amplitude and periodicity found in the O-C diagram. The mask is used to constrain the N-body retrieval (i.e. search space) which significantly reduces the number of simulations needed to find a solution. + +![](ttv_prior.png) + +The example code has recently been refactored from [Nbody-AI](https://github.com/pearsonkyle/Nbody-ai) and is still in development. Some retrievals may not work entirely... please be paitent and report any issues to our slack. -Match your Lomb-Scargle periodicity to this chart to determine the mass and period of a potential outer planet. A digitized version of this plot is coming soon... Occasionally, the O-C perturbation will experience an additional periodicity with only a single perturber, it is suspected to be due to precession of the inner planet's orbit. For now, just worry about the 1st order solution when matching to your observations. +If you make use of this code please cite: -![](grid_100_50.png) +https://ui.adsabs.harvard.edu/abs/2019AJ....158..243P/abstract diff --git a/examples/nbody/example.py b/examples/nbody/example.py new file mode 100644 index 00000000..64c1dd16 --- /dev/null +++ b/examples/nbody/example.py @@ -0,0 +1,143 @@ +import time +import numpy as np +import matplotlib.pyplot as plt +from astropy import units as u + +from exotic.api.nbody import report, generate, nbody_fitter, analyze, estimate_prior, TTV, interp_distribution + +mearth = u.M_earth.to(u.kg) +msun = u.M_sun.to(u.kg) +mjup = u.M_jup.to(u.kg) + +# create some sample data +objects = [ + # units: Msun, Days, au + {'m':0.95}, # stellar mass + {'m':1.169*mjup/msun, 'P':2.797436, 'inc':3.14159/2, 'e':0, 'omega':0 }, + {'m':0.1*mjup/msun, 'P':2.797436*1.9, 'inc':3.14159/2, 'e':0.0, 'omega':0 }, +] # HAT-P-37 + +# create REBOUND simulation +n_orbits = 2000 + +# time the simulation +t1 = time.time() +# inputs: object dict, length of simulation in days, number of timesteps [1hr] (should be at least 1/20 orbital period) +sim_data = generate(objects, objects[1]['P']*n_orbits, int(n_orbits*objects[1]['P']*24) ) +t2 = time.time() +print(f"Simulation time: {t2-t1:.2f} seconds") + +# collect the analytics of interest from the simulation +# lomb-scargle can be a lil slow +ttv_data = analyze(sim_data) + +# plot the results +report(ttv_data) + +# create a fake dataset +tmids = 2459150 + ttv_data['planets'][0]['tt'] +# add random noise to observations +tmids += np.random.normal(0,0.5,len(tmids))/(24*60) +# randomly select 50 observations without repeat +tmids = np.random.choice(tmids,50,replace=False) +# add random error to observations between +err = 1/24/60 + np.random.random()*0.25/24/60 + np.random.normal(0,0.1,len(tmids))/(24*60) +# estimate orbital epochs +orbit = np.rint((tmids-tmids.min())/ttv_data['planets'][0]['P']).astype(int) + +# estimate period from data +ttv,P,T0 = TTV(orbit, tmids) + +# run though linear fitter to estimate prior +from exotic.api.ephemeris import ephemeris_fitter + +# min and max values to search between for fitting +bounds = { + 'P': [P - 0.1, P + 0.1], # orbital period + 'T0': [T0 - 0.1, T0 + 0.1] # mid-transit time +} + +# used to plot red overlay in O-C figure +prior = { + 'P': [P, 0.00001], # value from linear lstq + 'T0': [T0, 0.001] # value from linear lstq +} + +lf = ephemeris_fitter(tmids, err, bounds, prior=prior) + +fig,ax = lf.plot_oc() +plt.tight_layout() +plt.show() + +# search for periodic signals in the data +fig,ax = lf.plot_periodogram() +plt.tight_layout() +plt.savefig('periodogram.png') +plt.show() +plt.close() + +# estimate ttv amplitude +amp = lf.amplitudes[0]*24*60 +per = lf.best_periods[0] # 1st order solution + +# estimate prior using periods from linear fit periodogram +masses, per_ratio, rvs, fig, ax = estimate_prior(amp, per) +masses *= mearth/msun +periods = per_ratio*lf.parameters['P'] +prior_fn_mass = interp_distribution(masses) +prior_fn_per = interp_distribution(periods) +plt.tight_layout() +plt.savefig('ttv_prior.png') +plt.show() +plt.close() + +# Parameters for N-body Retrieval +nbody_prior = [ + # star + {'m':0.95}, + + # inner planet + {'m':1.169*mjup/msun, + 'P':lf.parameters['P'], + 'inc':3.14159/2, + 'e':0, + 'omega':0}, + + # outer planet + {'m':masses.mean(), + 'P':periods.mean(), + 'inc':3.14159/2, + 'e':0, + 'omega':0,}, +] + +# specify data to fit +data = [ + {}, # data for star (e.g. RV) + {'Tc':tmids, 'Tc_err':err}, # data for inner planet (e.g. Mid-transit times) + {} # data for outer planet (e.g. Mid-transit times) +] + + +# set up where to look for a solution +nbody_bounds = [ + {}, # no bounds on stellar parameters + + { # bounds for inner planet + 'P': [nbody_prior[1]['P']-0.025, nbody_prior[1]['P']+0.025], # based on solution from linear fit\ + # 'T0': [None, None] # mid-transit is automatically adjusted, no need to include in bounds + }, + { # bounds for outer planet + 'P':[periods.min(), periods.max()], # period [day] + 'P_logl': prior_fn_per, # prior likelihood function + 'm':[masses.min(), masses.max()], # mass [msun + 'm_logl': prior_fn_mass, # prior likelihood function + 'omega': [-np.pi, np.pi] # argument of periastron [rad] + } +] + +# run the fitter +nfit = nbody_fitter(data, nbody_prior, nbody_bounds) + +# print(nfit.parameters) +# print(nfit.errors) diff --git a/examples/nbody/grid_100_50.png b/examples/nbody/grid_100_50.png deleted file mode 100644 index 512791b8..00000000 Binary files a/examples/nbody/grid_100_50.png and /dev/null differ diff --git a/examples/nbody/grid_500_150.png b/examples/nbody/grid_500_150.png new file mode 100644 index 00000000..7104c192 Binary files /dev/null and b/examples/nbody/grid_500_150.png differ diff --git a/exotic/api/ephemeris.py b/exotic/api/ephemeris.py index f9bff13b..954e6fde 100644 --- a/exotic/api/ephemeris.py +++ b/exotic/api/ephemeris.py @@ -41,26 +41,16 @@ # a periodogram analysis. # # ########################################################################### # -from astropy.timeseries import LombScargle -from astropy import units as u -import copy -from itertools import cycle -import matplotlib.pyplot as plt import numpy as np -import rebound +from itertools import cycle import statsmodels.api as sm +import matplotlib.pyplot as plt +from astropy.timeseries import LombScargle try: from ultranest import ReactiveNestedSampler except ImportError: - import dynesty - import dynesty.plotting - from dynesty.utils import resample_equal - from scipy.stats import gaussian_kde + print("Warning: ultranest not installed. Nested sampling will not work.") -try: - from nested_linear_fitter import linear_fitter, non_linear_fitter -except ImportError: - from .nested_linear_fitter import linear_fitter, non_linear_fitter try: from plotting import corner except ImportError: @@ -87,7 +77,10 @@ def __init__(self, data, dataerr, bounds=None, prior=None, labels=None, verbose= self.data = data self.dataerr = dataerr self.bounds = bounds - self.labels = np.array(labels) + if labels is not None: + self.labels = np.array(labels) + else: + self.labels = np.array(["Data"]*len(data)) self.prior = prior.copy() # dict {'P':(0.1,0.5), 'T0':(0,1), (value, uncertainty)} self.verbose = verbose @@ -173,7 +166,6 @@ def plot_oc(self, savefile=None, ylim='none', show_2sigma=False, prior_name="Pri for i, ulabel in enumerate(ulabels): # find where the label matches mask = self.labels == ulabel - # plot the data/residuals ax.errorbar(self.epochs[mask], self.residuals[mask] * 24 * 60, yerr=self.dataerr[mask] * 24 * 60, ls='none', marker=next(markers), color=next(colors), label=ulabel) else: @@ -340,7 +332,7 @@ def plot_triangle(self): ) return fig - def plot_periodogram(self, minper=0, maxper=0, minper2=1, maxper2=9): + def plot_periodogram(self, minper=0, maxper=0, minper2=0, maxper2=0): """ Search the residuals for periodic signals. """ ######################################## @@ -392,19 +384,30 @@ def plot_periodogram(self, minper=0, maxper=0, minper2=1, maxper2=9): ######################################## # subtract first order solution from data and recompute periodogram - maxper = maxper2 + if minper2 == 0: + minper2 = per*2.1 + if maxper2 == 0: + maxper2 = (np.max(self.epochs) - np.min(self.epochs)) * 2. + + # reset minper2 if it is greater than maxper2 + if minper2 > maxper2: + minper2 = 3.1 # recompute on new grid ls2 = LombScargle(self.epochs, residuals_linear, dy=self.dataerr) # search for periods greater than first order freq2,power2 = ls.autopower(maximum_frequency=1./(minper2+1), - minimum_frequency=1./maxper, nyquist_factor=2) + minimum_frequency=1./maxper2, nyquist_factor=2) # TODO do a better job defining period grid, ignoring harmonics of first order solution +- 0.25 day # find max period - mi2 = np.argmax(power2) + try: + mi2 = np.argmax(power2) + except: + import pdb; pdb.set_trace() + per2 = 1. / freq2[mi2] # create basis vectors for second order solution @@ -574,6 +577,7 @@ def plot_periodogram(self, minper=0, maxper=0, minper2=1, maxper2=9): # perform the weighted least squares regression to find second order fourier solution res = sm.WLS(residuals_first_order, basis2.T, weights=1.0 / self.dataerr ** 2).fit() coeffs = res.params # retrieve the slope and intercept of the fit from res + self.coeffs = res.params y_bestfit = np.dot(basis2.T, coeffs) # super sample fourier solution @@ -636,6 +640,11 @@ def plot_periodogram(self, minper=0, maxper=0, minper2=1, maxper2=9): ax[3].legend(loc='best') + self.best_periods = np.array([per, per2]) + coeff1 = np.sqrt(coeffs[1]**2 + coeffs[1]**2) + coeff2 = np.sqrt(coeffs[3]**2 + coeffs[4]**2) + self.amplitudes = np.array([coeff1, coeff2]) + return fig, ax class decay_fitter(object): @@ -659,7 +668,10 @@ def __init__(self, data, dataerr, bounds=None, prior=None, labels=None, verbose= self.data = data self.dataerr = dataerr self.bounds = bounds - self.labels = np.array(labels) + if labels is not None: + self.labels = np.array(labels) + else: + self.labels = ["Data"]*len(data) self.prior = prior.copy() # dict {'m':(0.1,0.5), 'b':(0,1)} self.verbose = verbose if bounds is None: @@ -924,7 +936,6 @@ def plot_triangle(self): ) return fig - if __name__ == "__main__": Tc = np.array([ # measured mid-transit times @@ -1030,7 +1041,7 @@ def plot_triangle(self): 0.000780,0.000610,0.001330,0.000770,0.000610,0.000520,0.001130,0.001130,0.000530,0.000780, 0.000420,0.001250,0.000380,0.000720,0.000860,0.000470,0.000950,0.000540]) - labels = np.full(len(Tc), 'EpW') + labels = np.full(len(Tc), 'Data') P = 1.0914203 # orbital period for your target @@ -1102,7 +1113,6 @@ def plot_triangle(self): plt.close() print("image saved to: periodogram.png") - # min and max values to search between for fitting bounds = { 'P': [P - 0.1, P + 0.1], # orbital period diff --git a/exotic/api/nbody.py b/exotic/api/nbody.py new file mode 100644 index 00000000..068cf0cd --- /dev/null +++ b/exotic/api/nbody.py @@ -0,0 +1,755 @@ +# ########################################################################### # +# Copyright (c) 2019-2020, California Institute of Technology. +# All rights reserved. Based on Government Sponsored Research under +# contracts NNN12AA01C, NAS7-1407 and/or NAS7-03001. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# 1. Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in +# the documentation and/or other materials provided with the +# distribution. +# 3. Neither the name of the California Institute of +# Technology (Caltech), its operating division the Jet Propulsion +# Laboratory (JPL), the National Aeronautics and Space +# Administration (NASA), nor the names of its contributors may be +# used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE CALIFORNIA +# INSTITUTE OF TECHNOLOGY BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED +# TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# ########################################################################### # +# EXOplanet Transit Interpretation Code (EXOTIC) +# # NOTE: See companion file version.py for version info. +# ########################################################################### # +# ########################################################################### # +# Code for running N-body simulations and computing transit timing variations. +# Adapted from: https://github.com/pearsonkyle/Nbody-ai +# ########################################################################### # +# ########################################################################### # +import os +import copy +import time +import numpy as np +import matplotlib.pyplot as plt +import rebound +from exotic.api.plotting import corner +from ultranest import ReactiveNestedSampler +from astropy.io import fits +from astropy import units as u +from scipy.interpolate import interp1d +from astropy.timeseries import LombScargle +from scipy.signal import find_peaks + +mearth = u.M_earth.to(u.kg) +msun = u.M_sun.to(u.kg) +mjup = u.M_jup.to(u.kg) + +# outlier robust "average" +maxavg = lambda x: (np.percentile(np.abs(x),75)+np.max(np.abs(x)) )*0.5 + +def empty_data(N): # orbit parameters in Nbody simulation + return { + 'x':np.zeros(N), + 'y':np.zeros(N), + 'z':np.zeros(N), + 'P':np.zeros(N), + 'a':np.zeros(N), + 'e':np.zeros(N), + 'inc':np.zeros(N), + 'Omega':np.zeros(N), + 'omega':np.zeros(N), + 'M':np.zeros(N), + } + +def generate(objects, Ndays=None, Noutputs=None): + # create rebound simulation + # for object parameters see: + # https://rebound.readthedocs.io/en/latest/_modules/rebound/particle.html + sim = rebound.Simulation() + sim.units = ('day', 'AU', 'Msun') + for i in range(len(objects)): + sim.add( **objects[i] ) + sim.move_to_com() + + if Ndays and Noutputs: + return integrate(sim, objects, Ndays, Noutputs) + else: + return sim + +def integrate(sim, objects, Ndays, Noutputs): + # ps is now an array of pointers and will change as the simulation runs + ps = sim.particles + + times = np.linspace(0., Ndays, Noutputs) # units of days + + pdata = [empty_data(Noutputs) for i in range(len(objects)-1)] + star = {'x':np.zeros(Noutputs),'y':np.zeros(Noutputs),'z':np.zeros(Noutputs) } + + # run simulation time steps + for i,time in enumerate(times): + sim.integrate(time) + + # record data + for k in star.keys(): + star[k][i] = getattr(ps[0], k) + + for j in range(1,len(objects)): + for k in pdata[j-1].keys(): + pdata[j-1][k][i] = getattr(ps[j],k) + + sim_data = ({ + 'pdata':pdata, + 'star':star, + 'times':times, + 'objects':objects, + 'dt': Noutputs/(Ndays*24*60*60) # conversion factor to get seconds for RV semi-amp + }) + + return sim_data + +def analyze(m, ttvfast=False): + + if ttvfast: + tt = transit_times( m['pdata'][0]['x'], m['star']['x'], m['times'] ) + ttv,per,t0 = TTV(np.arange(len(tt)),tt ) + return np.arange(len(ttv)),ttv,tt + + RV = np.diff(m['star']['x'])*1.496e11*m['dt'] # measurements -> seconds + + freq,power,peak_periods = lomb_scargle(m['times'][1:], RV, min_freq=1./365, max_freq=1.) + + data = { + 'times':m['times'], + 'RV':{'freq':freq, 'power':power, 'signal':RV, 'max':maxavg(RV), 'peak_periods':peak_periods}, + 'mstar': m['objects'][0]['m'], + 'planets':[], + 'objects':m['objects'] + } + + # parse planet data + for j in range(1,len(m['objects'])): + key = label='Planet {}'.format(j) + pdata = {} + + # compute transit times + tt = transit_times( m['pdata'][j-1]['x'], m['star']['x'], m['times'] ) + if len(tt)>=3: + ttv,per,t0 = TTV(np.arange(len(tt)),tt ) + else: + per = m['objects'][j]['P'] + t0 = 0 + ttv = [0] + + # periodogram of o-c + freq,power,peak_periods = lomb_scargle(np.arange(len(ttv)),ttv) + + # save data + for k in ['e','inc','a', 'omega']: + pdata[k] = np.mean( m['pdata'][j-1][k] ) + pdata['mass'] = m['objects'][j]['m'] + pdata['P'] = per + pdata['t0'] = t0 + pdata['tt'] = tt + pdata['ttv'] = ttv + pdata['max'] = maxavg(ttv) + pdata['freq'] = freq + pdata['power'] = power + pdata['peak_periods'] = peak_periods + + # down sample data for plotting + pdata['x'] = m['pdata'][j-1]['x'][::4] + pdata['y'] = m['pdata'][j-1]['z'][::4] + data['planets'].append(pdata) + + return data + +def lomb_scargle(t,y,dy=None, min_freq=None, max_freq=None, npeaks=0, peaktol=0.05): + + # perform quick periodicity analysis on residuals + find peaks in power plot + if dy is not None: + ls = LombScargle(t, y, dy) + else: + ls = LombScargle(t, y) + + if min_freq is None: + min_freq = 1./(1.5*(max(t)-min(t))) + + if max_freq is None: + max_freq = 1./2.5 + + # compute power spectrum + freq,power = ls.autopower(maximum_frequency=max_freq, minimum_frequency=min_freq, + nyquist_factor=2, samples_per_peak=5, method='cython') + + # find peaks in power spectrum + peaks,amps = find_peaks(power,height=peaktol) + + # sort by biggest to smallest amplitude + peaks = peaks[np.argsort(amps['peak_heights'])[::-1]] + peak_periods = 1./freq[peaks] + return freq,power,peak_periods + +def report(data, savefile=None): + + # set up simulation summary report + f = plt.figure( figsize=(12,8) ) + plt.subplots_adjust() + ax = [ plt.subplot2grid( (2,3), (0,0) ), # x,y plot + plt.subplot2grid( (2,3), (1,0) ), # table data + plt.subplot2grid( (2,3), (0,1) ), # # RV semi-amplitude plot #include theoretical limit for 2body system + plt.subplot2grid( (2,3), (1,1) ), # O-C plot # TODO add second axis for time (day) + plt.subplot2grid( (2,3), (0,2) ), # lomb scargle for RV semi-amplitude + plt.subplot2grid( (2,3), (1,2) ) # lomb scargle for o-c + ] + plt.subplots_adjust(top=0.96, bottom=0.07, left=0.1, right=0.98, hspace=0.3, wspace=0.3) + + ax[2].plot(data['times'][1:],data['RV']['signal'],'k-' ) + ax[2].set_xlim([0, 2.5*data['planets'][-1]['P'] ]) + ax[2].set_ylabel('RV semi-amplitude (m/s)') + ax[2].set_xlabel('time (day)') + + ax[4].plot(1./data['RV']['freq'],data['RV']['power'] ) + # plot reconstructed signal from sin series + + ax[4].set_xlabel('Period (day)') + ax[4].set_ylabel('|RV semi-amplitude| Power') + ax[4].set_xlim([1, 1.5*data['planets'][-1]['P'] ]) + #u = (m['objects'][1]['m']/m['objects'][0]['m']) * np.sqrt( G*(m['objects'][0]['m']+m['objects'][1]['m'])/m['objects'][1]['a'] )*1.496e11/(24*60*60) + #print('expected RV:',u) + + # create table stuff + keys = ['mass','a','P','inc','e','max'] + units = [msun/mearth,1,1,1,1,24*60] + rounds = [2,3,2,2,3,1,1] + + # make 2d list for table + tdata=[] + for i in range(len(units)): + tdata.append( [0 for i in range(len(data['planets']))] ) + + tinfo = { + 'mass':'Mass (Mearth)', + 'a':'Semi-major Axis (au)', + 'P':'Period (day)', + 'inc':'Inclination (deg)', + 'e':'Eccentricity', + 'max':'TTV Max Avg (min)', + } + row_labels = [ tinfo[k] for k in keys ] + col_labels = [ "P {}".format(i+1) for i in range(len(data['planets']))] + + # for each planet in the system + for j in range(1,len(data['objects'])): + + # plot orbits + ax[0].plot( data['planets'][j-1]['x'], data['planets'][j-1]['y'],label='Planet {}'.format(j),lw=0.5,alpha=0.5 ) + + if len(data['planets'][j-1]['ttv']) >= 2: + # plot O-C + ax[3].plot(data['planets'][j-1]['ttv']*24*60,label='Planet {}'.format(j) ) # convert days to minutes + + # lomb-scargle for O-C + ax[5].semilogx( 1./data['planets'][j-1]['freq'], data['planets'][j-1]['power'], label='Planet {}'.format(j) ) + + # add peak periods to legend + for i,per in enumerate(data['planets'][j-1]['peak_periods']): + # vertical line with label + if per > 0: + ax[5].axvline(per, ls='--', lw=0.5, label=f"{np.round(per,2)} day") + + # populate table data + for i,k in enumerate(keys): + tdata[i][j-1] = np.round( data['planets'][j-1][k] * units[i], rounds[i]) + + if k == 'inc': + tdata[i][j-1] = np.round( np.rad2deg(data['planets'][j-1][k]) * units[i], rounds[i]) + + table = ax[1].table(cellText=tdata, + colLabels=col_labels, + rowLabels=row_labels, + colWidths=[0.1]*len(col_labels), + loc='center') + + ax[1].set_title('Stellar mass: {:.2f} Msun'.format(data['mstar']) ) + table.scale(1.5,1.5) + + ax[1].axis('tight') + ax[1].axis('off') + + ax[0].set_ylabel('(au)') + ax[0].set_xlabel('(au)') + ax[0].legend(loc='best') + + ax[1].set_ylabel('position (au)') + ax[1].set_xlabel('time (day)') + + ax[3].set_ylabel('O-C (minutes)') + ax[3].set_xlabel('transit epoch') + ax[3].legend(loc='best') + ax[3].grid(True) + + ax[5].set_xlabel('Period (epoch)') + ax[5].set_ylabel('|O-C| Power' ) + #ax[5].set_xlim([1,30]) + ax[5].legend(loc='best') + + if savefile: + plt.savefig(savefile) + plt.close() + else: + plt.show() + +def find_zero(t1,x1, t2,x2): + """ + Find the zero of a line given two points. + + Parameters + ---------- + t1 : float + Time of first point. + + x1 : float + Position of first point. + + t2 : float + Time of second point. + + x2 : float + Position of second point. + + Returns + ------- + T0 : float + Time of zero. + """ + m = (x2-x1)/(t2-t1) + T0 = -x1/m + t1 + return T0 + +def transit_times(xp,xs,times): + """ + Find transit times from position data. + + Parameters + ---------- + xp : array + Planet x positions. + + xs : array + Star x positions. + + times : array + Simulation times. + + Returns + ------- + tt : array + Transit times. + """ + # check for sign change in position difference + dx = xp-xs + tt = [] + for i in range(1,len(dx)): + if dx[i-1] >= 0 and dx[i] <= 0: + tt.append( find_zero(times[i-1],dx[i-1], times[i],dx[i]) ) + return np.array(tt) + +def TTV(epochs, tt): + """ + Fit a line to the data and subtract it from the data. + + Parameters + ---------- + epochs : array + Orbits epochs of the data. + + tt : array + Transit times. + + Returns + ------- + ttv : array + Transit times with the linear trend subtracted. + + m : float + Slope of the linear trend. + + b : float + Intercept of the linear trend. + """ + N = len(epochs) + A = np.vstack([np.ones(N), epochs]).T + b, m = np.linalg.lstsq(A, tt, rcond=None)[0] + ttv = (tt-m*np.array(epochs)-b) + return [ttv,m,b] + + +# directory path of file +fpath = os.path.dirname(os.path.abspath(__file__)) +def estimate_prior( amplitude, periodicity, amp_lim=2, per_lim=4, file_name = os.path.join(fpath,"ttv_grid.fits"), show=True): + + with fits.open(file_name) as ttv_grid_list: + + # The grid is stored in multiple extensions + Npts = ttv_grid_list[0].header['NAXIS1'] + + # in Earth + mass_grid = np.linspace(ttv_grid_list[0].header['CRVAL1'], ttv_grid_list[0].header['CRVAL1'] + Npts*ttv_grid_list[0].header['CDELT1'], Npts) + mass_grid_2 = np.linspace(ttv_grid_list[0].header['CRVAL1'], ttv_grid_list[0].header['CRVAL1'] + Npts*ttv_grid_list[0].header['CDELT1'], int(0.5*Npts)) + period_grid = np.linspace(ttv_grid_list[0].header['CRVAL2'], ttv_grid_list[0].header['CRVAL2'] + Npts*ttv_grid_list[0].header['CDELT2'], Npts) + #period_grid_2 = np.linspace(ttv_grid_list[0].header['CRVAL2'], ttv_grid_list[0].header['CRVAL2'] + Npts*ttv_grid_list[0].header['CDELT2'], int(0.5*Npts)) + + # extract data + pg, mg = np.meshgrid(period_grid, mass_grid) + amplitude_grid = ttv_grid_list[0].data + periodicity1_grid = ttv_grid_list[1].data + periodicity2_grid = ttv_grid_list[2].data + # combine periods by min + #periodicity_grid = np.minimum(periodicity1_grid, periodicity2_grid) + rv_grid = ttv_grid_list[3].data + + # amplitude grid + fig, ax = plt.subplots(3, 2, figsize=(9, 11)) + fig.suptitle("N-body Estimates from Transit Timing Variations", fontsize=16) + plt.subplots_adjust(left=0.0485, right=0.985, bottom=0.15, top=0.9, wspace=0.3) + im = ax[0,0].imshow(amplitude_grid, origin='lower', extent=[mass_grid[0], mass_grid[-1],period_grid[0], period_grid[-1]], vmin=0,vmax=30, aspect='auto', cmap='jet', interpolation='none') + ax[0,0].set_ylabel('Period Ratio') + ax[0,0].set_xlabel('Mass [Earth]') + cbar = fig.colorbar(im, ax=ax[0,0]) + cbar.set_label('TTV Amplitude [min]') + + # periodoicity 1 grid + im = ax[0,1].imshow(periodicity1_grid, origin='lower', extent=[mass_grid[0], mass_grid[-1],period_grid[0], period_grid[-1]], vmin=0,vmax=0.5*np.nanmax(ttv_grid_list[1].data), aspect='auto', cmap='jet', interpolation='none') + ax[0,1].set_ylabel('Period Ratio') + ax[0,1].set_xlabel('Mass [Earth]') + cbar = fig.colorbar(im, ax=ax[0,1]) + cbar.set_label('TTV Periodicity 1st Order [epoch]') + + # simulate some results based on the periodogram + mask = (amplitude_grid > amplitude-amp_lim) & (amplitude_grid < amplitude+amp_lim) & \ + (((periodicity1_grid > periodicity-per_lim) & (periodicity1_grid < periodicity+per_lim)) | \ + ((periodicity2_grid > periodicity-per_lim) & (periodicity2_grid < periodicity+per_lim))) + # plot mask + im = ax[1,0].imshow(mask, origin='lower', extent=[mass_grid[0], mass_grid[-1],period_grid[0], period_grid[-1]], vmin=0,vmax=1, aspect='auto', cmap='binary_r', interpolation='none') + ax[1,0].set_ylabel('Period Ratio') + ax[1,0].set_xlabel('Mass [Earth]') + cbar = fig.colorbar(im, ax=ax[1,0]) + cbar.set_label('N-body Prior from O-C Data') + + # compare to original + # TODO find why the last bin has double the number of points + masses = mg.T[mask].flatten() + ax[1,1].hist(masses, bins=mass_grid_2, alpha=0.5) + ax[1,1].set_xlabel('Mass [Earth]') + ax[1,1].set_ylabel('Prior Probability') + ax[1,1].axes.yaxis.set_ticklabels([]) + ax[1,1].set_xlim([mg.min(), mg.max()-10]) + + # compare to original + periods = pg.T[mask].flatten() + ax[2,0].hist(periods, bins=period_grid, density=True, alpha=0.5) + ax[2,0].set_xlabel('Period Ratio') + ax[2,0].set_ylabel('Prior Probability') + ax[2,0].axes.yaxis.set_ticklabels([]) + ax[2,0].set_xlim([pg.min(), pg.max()]) + + # plot histogram for rv data + rvs = rv_grid.T[mask].flatten() + ax[2,1].hist(rvs, bins=50, density=True, alpha=0.5) + ax[2,1].set_xlabel('RV Semi-Amplitude [m/s]') + ax[2,1].set_ylabel('Prior Probability') + ax[2,1].axes.yaxis.set_ticklabels([]) + + # m_earth, days, m/s + return masses, periods, rvs, fig, ax + +def interp_distribution(values, nbins=50): + value_grid = np.linspace(np.min(values), np.max(values), nbins) + heights, edges = np.histogram(values, bins=value_grid, density=True) + edge_center = (edges[:-1] + edges[1:]) / 2 + + # Step 2: Normalize the histogram + bin_widths = np.diff(edges) + total_area = np.sum(bin_widths * heights) + normalized_heights = heights / total_area + + # Step 3: Interpolate to create a continuous PDF + return interp1d(edge_center, normalized_heights, kind='linear', bounds_error=False, fill_value=0) + + + + +class nbody_fitter(): + + def __init__(self, data, prior=None, bounds=None, verbose=True): + self.data = data + self.bounds = bounds + self.prior = prior + self.verbose = verbose + self.fit_nested() + + def fit_nested(self): + + # set up some arrays for mapping sampler output + freekeys = [] + boundarray = [] + for i,planet in enumerate(self.bounds): + for bound in planet: + if '_logl' in bound or '_fn' in bound: + continue + freekeys.append(f"{i}_{bound}") + boundarray.append(planet[bound]) + + # find min and max time for simulation + min_time = np.min(self.data[1]['Tc']) + max_time = np.max(self.data[1]['Tc']) + # todo extend for multiplanet systems + + sim_time = (max_time-min_time)*1.05 + # TODO extend for multiplanet systems + Tc_norm = self.data[1]['Tc'] - min_time # normalize the data to the first observation + self.orbit = np.rint(Tc_norm / self.prior[1]['P']).astype(int) # number of orbits since first observation (rounded to nearest integer) + + # numpify + boundarray = np.array(boundarray) + bounddiff = np.diff(boundarray,1).reshape(-1) + + # create queue and save simulation to + def loglike(pars): + chi2 = 0 + + # set parameters + for i,par in enumerate(pars): + idx,key = freekeys[i].split('_') + idx = int(idx) + if key == 'tmid': + continue + # this dict goes to REBOUND and needs specific keys + self.prior[idx][key] = par + + # prior likelihood function for P, m of outer planet + likelihood_m = self.bounds[2]['m_logl'](self.prior[2]['m']) + likelihood_P = self.bounds[2]['P_logl'](self.prior[2]['P']) + if likelihood_m <= 0 or likelihood_P <= 0: + return -1e6 + + # take difference between data and simulation + # run N-body simulation + sim_data = generate(self.prior, sim_time, int(sim_time*24)) # uses linspace behind the scenes + + # json structure with analytics of interest from the simulation + # ttv_data = analyze(sim_data) # slow + + sim_shift = 0 + # loop over planets, check for data + for i,planet in enumerate(self.prior): + if self.data[i]: + # compute transit times from N-body simulation + Tc_sim = transit_times( sim_data['pdata'][i-1]['x'], sim_data['star']['x'], sim_data['times'] ) + + # derive an offset in time from the first planet + if i-1==0: + sim_shift = Tc_sim.min() + + # shift the first mid-transit in simulation to observation + Tc_sim -= sim_shift # first orbit has tmid at 0 + + # scale Tc_sim to data + try: + residual = self.data[i]['Tc'] - Tc_sim[self.orbit] + except: + #import pdb; pdb.set_trace + #ttv,m,b = TTV(np.arange(Tc_sim.shape[0]), Tc_sim) + #ttv1,m1,b1 = TTV(self.orbit, self.data[i]['Tc']) + # could be unstable orbit or not enough data + # switch to average error and clip by max epoch? + # print(self.prior) + chi2 += -1e6 + continue + + import pdb; pdb.set_trace + Tc_sim += residual.mean() # why? + + # take difference between data and simulation + try: + chi2 += -0.5*np.sum(((self.data[i]['Tc'] - Tc_sim[self.orbit])/self.data[i]['Tc_err'])**2) + except: + chi2 += -1e6 + # print(self.prior) + # usually unstable orbit + + return chi2 + + + def prior_transform(upars): + return (boundarray[:,0] + bounddiff*upars) + + if self.verbose: + self.results = ReactiveNestedSampler(freekeys, loglike, prior_transform).run(max_ncalls=1e5) + else: + self.results = ReactiveNestedSampler(freekeys, loglike, prior_transform).run(max_ncalls=1e5, show_status=self.verbose, +viz_callback=self.verbose) + + self.errors = {} + self.quantiles = {} + self.parameters = copy.deepcopy(self.prior) + + # TODO finish saving final results + posteriors + + #for i, key in enumerate(freekeys): + # self.parameters[key] = self.results['maximum_likelihood']['point'][i] + # self.errors[key] = self.results['posterior']['stdev'][i] + # self.quantiles[key] = [ + # self.results['posterior']['errlo'][i], + # self.results['posterior']['errup'][i]] + + +if __name__ == '__main__': + + # create some sample data + objects = [ + # units: Msun, Days, au + {'m':0.95}, # stellar mass + {'m':1.169*mjup/msun, 'P':2.797436, 'inc':3.14159/2, 'e':0, 'omega':0 }, + {'m':0.1*mjup/msun, 'P':2.797436*1.9, 'inc':3.14159/2, 'e':0.0, 'omega':0 }, + ] # HAT-P-37 + + # create REBOUND simulation + n_orbits = 2000 + + # time the simulation + t1 = time.time() + # inputs: object dict, length of simulation in days, number of timesteps [1hr] (should be at least 1/20 orbital period) + sim_data = generate(objects, objects[1]['P']*n_orbits, int(n_orbits*objects[1]['P']*24) ) + t2 = time.time() + print(f"Simulation time: {t2-t1:.2f} seconds") + + # collect the analytics of interest from the simulation + # lomb-scargle can be a lil slow + ttv_data = analyze(sim_data) + + # plot the results + report(ttv_data) + + # create a fake dataset + tmids = 2459150 + ttv_data['planets'][0]['tt'] + # add random noise to observations + tmids += np.random.normal(0,0.5,len(tmids))/(24*60) + # randomly select 50 observations without repeat + tmids = np.random.choice(tmids,50,replace=False) + # add random error to observations between + err = 1/24/60 + np.random.random()*0.25/24/60 + np.random.normal(0,0.1,len(tmids))/(24*60) + # estimate orbital epochs + orbit = np.rint((tmids-tmids.min())/ttv_data['planets'][0]['P']).astype(int) + + # estimate period from data + ttv,P,T0 = TTV(orbit, tmids) + + # run though linear fitter to estimate prior + from exotic.api.ephemeris import ephemeris_fitter + + # min and max values to search between for fitting + bounds = { + 'P': [P - 0.1, P + 0.1], # orbital period + 'T0': [T0 - 0.1, T0 + 0.1] # mid-transit time + } + + # used to plot red overlay in O-C figure + prior = { + 'P': [P, 0.00001], # value from linear lstq + 'T0': [T0, 0.001] # value from linear lstq + } + + lf = ephemeris_fitter(tmids, err, bounds, prior=prior) + + fig,ax = lf.plot_oc() + plt.tight_layout() + plt.show() + + # search for periodic signals in the data + fig,ax = lf.plot_periodogram() + plt.tight_layout() + plt.savefig('periodogram.png') + plt.show() + plt.close() + + # estimate ttv amplitude + amp = lf.amplitudes[0]*24*60 + per = lf.best_periods[0] # 1st order solution + + # estimate prior using periods from linear fit periodogram + masses, per_ratio, rvs, fig, ax = estimate_prior(amp, per) + masses *= mearth/msun + periods = per_ratio*lf.parameters['P'] + prior_fn_mass = interp_distribution(masses) + prior_fn_per = interp_distribution(periods) + plt.tight_layout() + plt.savefig('ttv_prior.png') + plt.show() + plt.close() + + # Parameters for N-body Retrieval + nbody_prior = [ + # star + {'m':0.95}, + + # inner planet + {'m':1.169*mjup/msun, + 'P':lf.parameters['P'], + 'inc':3.14159/2, + 'e':0, + 'omega':0}, + + # outer planet + {'m':masses.mean(), + 'P':periods.mean(), + 'inc':3.14159/2, + 'e':0, + 'omega':0,}, + ] + + # specify data to fit + data = [ + {}, # data for star (e.g. RV) + {'Tc':tmids, 'Tc_err':err}, # data for inner planet (e.g. Mid-transit times) + {} # data for outer planet (e.g. Mid-transit times) + ] + + # TODO break P,m modes into individual runs + + # set up where to look for a solution + nbody_bounds = [ + {}, # no bounds on stellar parameters + + { # bounds for inner planet + 'P': [nbody_prior[1]['P']-0.025, nbody_prior[1]['P']+0.025], # based on solution from linear fit\ + # 'T0': [None, None] # mid-transit is automatically adjusted, no need to include in bounds + }, + { # bounds for outer planet + 'P':[periods.min(), periods.max()], # period [day] + 'P_logl': prior_fn_per, # prior likelihood function + 'm':[masses.min(), masses.max()], # mass [msun + 'm_logl': prior_fn_mass, # prior likelihood function + 'omega': [-np.pi, np.pi] # argument of periastron [rad] + } + ] + + # run the fitter + nfit = nbody_fitter(data, nbody_prior, nbody_bounds) + + # print(nfit.parameters) + # print(nfit.errors) \ No newline at end of file diff --git a/exotic/api/nested_linear_fitter.py b/exotic/api/nested_linear_fitter.py index 229d4c8a..c7890018 100644 --- a/exotic/api/nested_linear_fitter.py +++ b/exotic/api/nested_linear_fitter.py @@ -40,22 +40,15 @@ # a periodogram analysis. # # ########################################################################### # -from astropy import constants as const -from astropy.timeseries import LombScargle -from astropy import units as u -from collections import OrderedDict -import copy -from itertools import cycle -import matplotlib.pyplot as plt import numpy as np +from itertools import cycle import statsmodels.api as sm +import matplotlib.pyplot as plt +from astropy.timeseries import LombScargle try: from ultranest import ReactiveNestedSampler except ImportError: - import dynesty - import dynesty.plotting - from dynesty.utils import resample_equal - from scipy.stats import gaussian_kde + print("Warning: ultranest not installed. Nested sampling will not work.") try: from plotting import corner @@ -911,256 +904,3 @@ def plot_triangle(self): hist_kwargs={'color': 'black', } ) return fig - - -#def main(): -if __name__ == "__main__": - - Tc = np.array([ # measured mid-transit times - #Exoplanet Watch transit times array -2457025.7867,2457347.7655,2457694.8263,2458112.8495,2458492.6454,2459532.7834,2459604.81036,2459604.8137,2459614.6365,2459616.8209, -2459651.7466,2459903.8588,2459914.7783,2459915.8684,2459925.6921,2459939.8793,2459949.7047,2459959.5249,2459962.8037,2459973.7129, -2459974.798,2459986.8057,2459994.4489,2460009.7312,2458867.01587,2459501.1295, - -#ExoClock transit times array -2454508.977,2454515.525,2454801.477,2454840.769,2455123.447,2455140.91,2455147.459,2455147.459,2455148.551,2455148.552,2455158.372, -2455158.373,2455159.464,2455159.467,2455160.556,2455163.831,2455172.562,2455192.205,2455209.669,2455210.762,2455215.129,2455230.407, -2455238.046,2455254.419,2455265.331,2455494.53,2455494.531,2455509.81,2455510.902,2455530.547,2455542.553,2455548.011,2455565.472, -2455566.563,2455575.296,2455575.297,2455578.569,2455589.483,2455590.576,2455598.216,2455600.397,2455600.398,2455600.398,2455601.49, -2455601.49,2455601.49,2455603.673,2455612.404,2455623.318,2455623.319,2455624.41,2455624.41,2455663.702,2455842.693,2455842.694, -2455876.528,2455887.442,2455887.442,2455887.442,2455888.533,2455888.534,2455888.534,2455889.625,2455890.716,2455901.63,2455903.814, -2455903.814,2455920.185,2455923.459,2455926.733,2455939.829,2455945.287,2455946.378,2455946.379,2455947.47,2455948.561,2455951.835, -2455952.927,2455959.475,2455960.567,2455970.39,2455971.481,2455981.304,2455982.395,2455983.487,2455984.578,2455985.67,2455996.584, -2456000.95,2456005.315,2456006.406,2456014.046,2456175.577,2456245.427,2456249.794,2456273.805,2456282.536,2456284.719,2456297.816, -2456302.182,2456305.455,2456319.644,2456328.376,2456329.467,2456604.505,2456605.596,2456606.688,2456607.779,2456629.607,2456630.699, -2456654.71,2456659.076,2456662.35,2456663.441,2456664.533,2456674.356,2456677.63,2456688.544,2456694.002,2456703.824,2456711.464, -2456719.104,2456721.287,2456722.378,2456986.502,2457010.513,2457012.696,2457033.434,2457045.438,2457046.53,2457059.627,2457060.718, -2457067.267,2457068.358,2457103.284,2457345.579,2457368.5,2457378.321,2457390.327,2457391.418,2457426.343,2457426.344,2457427.435, -2457451.446,2457474.364,2457486.372,2457671.913,2457691.559,2457703.564,2457706.838,2457726.484,2457727.575,2457760.318,2457765.775, -2457766.866,2457772.324,2457773.415,2457776.689,2457786.512,2457788.695,2457800.7,2457808.34,2457809.432,2457809.434,2457810.523, -2457832.348,2457832.351,2458026.624,2458050.635,2458060.459,2458072.462,2458073.555,2458074.647,2458077.921,2458123.76,2458124.852, -2458132.491,2458134.675,2458136.858,2458155.41,2458155.412,2458156.503,2458159.778,2458166.326,2458178.331,2458214.347,2458411.895, -2458467.557,2458471.923,2458478.472,2458479.564,2458490.477,2458494.843,2458501.39,2458501.392,2458506.848,2458536.315,2458537.408, -2458871.382,2458880.113,2458883.384,2458893.209,2458903.032,2459088.575,2459148.602,2459156.242,2459157.334,2459190.076,2459194.441, -2459194.443,2459196.623,2459197.717,2459204.264, - -#ETD mid transit times array -2454508.976854,2454836.403393,2454840.768603,2454840.771193,2454908.437992,2454931.358182,2455164.923956,2455164.924316,2455172.561816,2455172.562786, -2455198.756735,2455230.406730,2455254.418870,2455256.596033,2455265.334053,2455506.533285,2455509.808915,2455519.632074,2455532.729964,2455541.461084, -2455576.387242,2455577.475172,2455577.476782,2455577.477112,2455578.568422,2455593.850192,2455600.398282,2455601.489211,2455603.671161,2455612.400801, -2455612.401321,2455615.675911,2455647.328760,2455857.973783,2455857.973783,2455867.797832,2455873.253612,2455877.618482,2455877.619162,2455889.623302, -2455899.448441,2455899.451251,2455900.539781,2455904.905271,2455905.995241,2455912.544561,2455912.544561,2455913.636131,2455914.726741,2455922.370121, -2455928.913321,2455936.552560,2455936.561520,2455937.645430,2455938.738730,2455946.383030,2455947.468480,2455949.650900,2455949.650900,2455950.742510, -2455957.287600,2455957.294400,2455957.296780,2455957.296780,2455961.659130,2455962.749930,2455970.392369,2455982.394919,2455993.309869,2455993.313119, -2455994.398719,2455994.399359,2456003.129689,2456003.130169,2456005.315328,2456222.509695,2456246.516835,2456249.793845,2456269.425195,2456271.623265, -2456296.724384,2456304.364534,2456304.365804,2456304.367454,2456304.367784,2456329.469294,2456364.392544,2456584.860173,2456595.772813,2456596.866383, -2456604.503073,2456608.869564,2456616.511354,2456626.333324,2456627.429034,2456628.515004,2456628.517074,2456630.699754,2456639.430344,2456639.432894, -2456640.522554,2456641.610274,2456650.341584,2456663.443614,2456688.544734,2456722.372535,2456722.382615,2456734.383505,2456734.388675,2456746.385515, -2456746.391885,2456927.546138,2456950.485119,2456960.306869,2456963.578439,2456963.578959,2456963.582479,2456986.501140,2456998.508180,2456998.508180, -2457032.333311,2457033.433561,2457069.448912,2457080.367572,2457080.367572,2457092.368353,2457344.486723,2457344.488693,2457345.578813,2457345.579833, -2457356.493733,2457357.585013,2457357.586493,2457368.498964,2457368.499654,2457369.589474,2457379.414424,2457379.414424,2457414.334726,2457426.343356, -2457451.447237,2457474.364078,2457474.364348,2457486.371248,2457668.637786,2457726.484138,2457749.405049,2457750.494949,2457751.585649,2457751.586939, -2457757.043281,2457758.134851,2457760.317471,2457760.318171,2457761.408951,2457770.140361,2457772.327372,2457773.415162,2457774.506522,2457809.431273, -2457809.435043,2457832.350874,2458087.745843,2458096.474134,2458109.571734,2458131.400315,2458396.615162,2458396.617562,2458411.896052,2458455.554593, -2458457.734253,2458479.568743,2458489.383454,2458489.387804,2458490.477904,2458490.482824,2458501.391674,2458501.392014,2458501.396544,2458513.402254, -2458537.407634,2458538.507734,2458837.550156,2458848.461686,2458859.375956,2458860.465876,2458860.469616,2458871.379136,2458871.382556,2458872.470796, -2458883.387165,2458883.389555,2458884.478175,2458884.478425,2458884.479455,2458906.306925,2458930.332215,2459147.510251,2459171.522921,2459172.608741, -2459172.617851,2459183.531910,2459193.352480,2459194.442340,2459194.443610,2459196.623950,2459197.719250,2459220.632699,2459242.461499,2459265.384258, -2459265.384318,2459505.495609,2459553.514627,2459553.515137,2459564.427746,2459565.520376,2459566.615006,2459576.438096]) - - Tc_error = np.array([ - 0.0039,0.0058,0.0054,0.0022,0.0065,0.0024,0.00091,0.00097,0.0022,0.0018, -0.002,0.0021,0.0065,0.0036,0.0071,0.0031,0.0044,0.0021,0.0028,0.0023, -0.0039,0.0071,0.0015,0.0055,0.00092,0.00092, - -#ExoClock error array -0.0002,0.00013,0.00041,0.00047,0.0007,0.000417,0.0009,0.00042,0.001,0.0013,0.00057, -0.0013,0.00095,0.00061,0.0007,0.000324,0.00044,0.0005,0.000463,0.000405,0.00082,0.00011, -0.00066,0.00014,0.0016,0.00048,0.00063,0.00037,0.000313,0.00077,0.00029,0.00024,0.0011, -0.00022,0.00063,0.00079,0.0008,0.00094,0.001,0.00036,0.00083,0.00026,0.00078,0.000289, -0.00017,0.00039,0.001,0.00088,0.0009,0.00054,0.00086,0.0014,0.00063,0.00093,0.0013, -0.0003,0.00038,0.00062,0.00075,0.00059,0.0003,0.00038,0.0009,0.00023,0.00083,0.00093, -0.000324,0.00046,0.0002,0.00092,0.0019,0.00079,0.00036,0.00034,0.00022,0.00028,0.00011, -0.0001,0.00018,0.0004,0.00029,0.00029,0.00094,0.00047,0.00029,0.000324,0.000417,0.00037, -0.0004,0.00038,0.0004,0.00023,0.00033,0.00033,0.000394,0.000301,0.0003,0.000301,0.000301, -0.00046,0.00026,0.000382,0.00027,0.00029,0.0002,0.0003,0.00034,0.000706,0.00019,0.00043, -0.000336,0.00034,0.00019,0.00019,0.00032,0.00028,0.000324,0.00041,0.00029,0.00029,0.00026, -0.00034,0.00034,0.00046,0.00043,0.00039,0.000486,0.0005,0.00049,0.00049,0.000347,0.000359, -0.00022,0.00021,0.0003,0.00042,0.0004,0.0013,0.00034,0.00033,0.00055,0.0006,0.00023,0.00021, -0.0007,0.0013,0.00035,0.00025,0.00034,0.00037,0.00028,0.00023,0.0006,0.00028,0.00039, -0.00024,0.00022,0.00029,0.00026,0.00048,0.00032,0.0004,0.00018,0.0009,0.00021,0.0006, -0.0006,0.00056,0.00023,0.0003,0.0003,0.00022,0.00034,0.00028,0.00027,0.00035,0.00031, -0.00032,0.00033,0.0005,0.00031,0.00032,0.00091,0.00034,0.00038,0.0017,0.0004,0.0005, -0.00026,0.0006,0.0006,0.0008,0.0003,0.0009,0.0003,0.00044,0.0008,0.0007,0.0009, -0.0003,0.0007,0.0005,0.0003,0.0013,0.0007,0.0003,0.0004,0.0003,0.0012,0.0006, -0.0005,0.0013,0.0004, - -#ETD error array -0.000200,0.000600,0.000470,0.001000,0.001000,0.000980,0.001490,0.001290,0.000440,0.000140, -0.001410,0.000110,0.000140,0.000590,0.001290,0.001120,0.000870,0.000700,0.000350,0.000500, -0.001350,0.000380,0.000690,0.000850,0.000820,0.000470,0.000420,0.001090,0.000940,0.000830, -0.000780,0.000540,0.001730,0.000710,0.000710,0.000750,0.001260,0.000920,0.001290,0.000730, -0.000630,0.000570,0.000380,0.001030,0.001350,0.000570,0.000570,0.000780,0.000470,0.000900, -0.000730,0.000910,0.001230,0.001190,0.000720,0.000770,0.001020,0.000590,0.000590,0.000660, -0.001100,0.001040,0.000570,0.000570,0.001070,0.001320,0.000860,0.001160,0.000600,0.000760, -0.000680,0.000760,0.000630,0.000600,0.000440,0.000810,0.000740,0.000670,0.000900,0.000550, -0.000520,0.001460,0.000890,0.001560,0.000580,0.001640,0.001170,0.000510,0.000960,0.000510, -0.000920,0.000710,0.000900,0.000510,0.001050,0.000970,0.000880,0.000440,0.000740,0.000680, -0.000820,0.000800,0.001040,0.000760,0.000540,0.000890,0.001080,0.001120,0.000730,0.001390, -0.001410,0.001640,0.000740,0.000570,0.000930,0.000890,0.000610,0.000510,0.001450,0.001450, -0.001130,0.000460,0.000510,0.001190,0.001190,0.000600,0.000650,0.001070,0.001090,0.000490, -0.000610,0.000520,0.000430,0.000580,0.000460,0.000340,0.000890,0.000890,0.001310,0.000650, -0.000860,0.000770,0.000550,0.001350,0.000820,0.000550,0.000840,0.000530,0.000940,0.000720, -0.000370,0.000520,0.000670,0.001030,0.000630,0.000330,0.000990,0.000550,0.000620,0.000720, -0.000940,0.000580,0.000350,0.000570,0.001380,0.000640,0.001180,0.000700,0.000700,0.000710, -0.000550,0.000580,0.000680,0.001030,0.000860,0.000850,0.000300,0.000760,0.000680,0.000820, -0.000610,0.001450,0.000730,0.000700,0.001350,0.000820,0.000670,0.000940,0.000490,0.001130, -0.000540,0.000540,0.000700,0.000790,0.000840,0.000600,0.000520,0.000730,0.000640,0.001020, -0.000780,0.000610,0.001330,0.000770,0.000610,0.000520,0.001130,0.001130,0.000530,0.000780, -0.000420,0.001250,0.000380,0.000720,0.000860,0.000470,0.000950,0.000540]) - - labels = np.full(len(Tc), 'EpW') - # # Tc = np.array([ # measured mid-transit times - # # 2459150.837905, 2459524.851045, - # # 2459546.613126, 2459565.643663, 2459584.690470, 2459584.686476, - # # 2459909.739104, 2459957.337739, 2459169.880602, 2458416.424861, - # # 2458428.664794, 2459145.400124, 2458430.025044, 2459164.440695, - # # 2458415.064846, 2458435.465269, 2458412.344658, 2459150.840070, - # # 2459160.360691, 2458431.384897, 2459146.760284, 2459154.920516, - # # 2458417.784945, 2459161.720466, 2459167.160761, 2458427.304939, - # # 2458413.705181, 2459163.080605, 2459153.560229, 2459168.520861, - # # 2458425.945068, 2459148.120269, 2458434.105274, 2458432.745423, - # # 2459152.200558, 2459165.800918, 2459159.000645, 2459149.480289, - # # 2455870.450027, 2456663.347570, 2457420.884390, 2457658.888900]) - - # # Tc_error = np.array([ - # # 0.001674, 0.000715, - # # 0.000758, 0.001560, 0.000371, 0.001651, - # # 0.000955, 0.000176, 0.000154, 0.000357, - # # 0.000381, 0.000141, 0.000362, 0.000149, - # # 0.000336, 0.000368, 0.000379, 0.000153, - # # 0.000153, 0.000349, 0.000149, 0.000146, - # # 0.000385, 0.000146, 0.000153, 0.000360, - # # 0.000356, 0.000147, 0.000147, 0.000146, - # # 0.000363, 0.000142, 0.000357, 0.000368, - # # 0.000160, 0.000160, 0.000151, 0.000160, - # # 0.000140, 0.000120, 0.000800, 0.000140]) - - # # labels for a legend - # labels = np.array(['TESS', 'TESS', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar', - # 'TESS', 'EpW', 'ExoClock', 'Unistellar']) - - P = 1.0914203 # orbital period for your target - - Tc_norm = Tc - Tc.min() # normalize the data to the first observation - # print(Tc_norm) - orbit = np.rint(Tc_norm / P) # number of orbits since first observation (rounded to nearest integer) - # print(orbit) - - # make a n x 2 matrix with 1's in the first column and values of orbit in the second - A = np.vstack([np.ones(len(Tc)), orbit]).T - - # perform the weighted least squares regression - res = sm.WLS(Tc, A, weights=1.0 / Tc_error ** 2).fit() - # use sm.WLS for weighted LS, sm.OLS for ordinary LS, or sm.GLS for general LS - - params = res.params # retrieve the slope and intercept of the fit from res - std_dev = np.sqrt(np.diagonal(res.normalized_cov_params)) - - slope = params[1] - slope_std_dev = std_dev[1] - intercept = params[0] - intercept_std_dev = std_dev[0] - - # 3 sigma clip based on residuals - calculated = orbit * slope + intercept - residuals = (Tc - calculated) / Tc_error - mask = np.abs(residuals) < 3 - Tc = Tc[mask] - Tc_error = Tc_error[mask] - labels = labels[mask] - - # print(res.summary()) - # print("Params =",params) - # print("Error matrix =",res.normalized_cov_params) - # print("Standard Deviations =",std_dev) - - print("Weighted Linear Least Squares Solution") - print("T0 =", intercept, "+-", intercept_std_dev) - print("P =", slope, "+-", slope_std_dev) - - # min and max values to search between for fitting - bounds = { - 'm': [P - 0.1, P + 0.1], # orbital period - 'b': [intercept - 0.1, intercept + 0.1] # mid-transit time - } - - # used to plot red overlay in O-C figure - prior = { - 'm': [slope, slope_std_dev], # value from WLS (replace with literature value) - 'b': [intercept, intercept_std_dev] # value from WLS (replace with literature value) - } - - lf = linear_fitter(Tc, Tc_error, bounds, prior=prior, labels=labels) - - lf.plot_triangle() - plt.subplots_adjust(top=0.9, hspace=0.2, wspace=0.2) - plt.savefig("posterior.png") - plt.close() - print("image saved to: posterior.png") - - fig, ax = lf.plot_oc() - plt.tight_layout() - plt.savefig("oc.png") - plt.show() - plt.close() - print("image saved to: oc.png") - - fig, ax = lf.plot_periodogram() - plt.tight_layout() - plt.savefig("periodogram.png") - plt.close() - print("image saved to: periodogram.png") - - - # min and max values to search between for fitting - bounds = { - 'P': [P - 0.1, P + 0.1], # orbital period - 't0': [intercept - 0.1, intercept + 0.1], # mid-transit time - 'dPdN': [-0.00001, 0.00001] # dPdN - } - - # used to plot red overlay in O-C figure - prior = { - 'P': [slope, slope_std_dev], # value from WLS (replace with literature value) - 't0': [intercept, intercept_std_dev], # value from WLS (replace with literature value) - 'dPdN': [0, 0.0001] # best guess - } - - nlf = non_linear_fitter(Tc, Tc_error, bounds, prior=prior, labels=labels) - - nlf.plot_triangle() - plt.subplots_adjust(top=0.9, hspace=0.2, wspace=0.2) - plt.savefig("nl_posterior.png") - plt.close() - print("image saved to: nl_posterior.png") - - nfig, ax = nlf.plot_oc() - plt.tight_layout() - plt.savefig("nl_oc.png") - plt.show() - plt.close() - print("image saved to: nl_oc.png") - - # for each free key print parameter and error - for key in nlf.parameters: - print(f"Parameter {key} = {nlf.parameters[key]:.2e} +- {nlf.errors[key]:.2e}") \ No newline at end of file