Merge pull request #1234 from rzellem/nbody

New example + API clean up

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

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)

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