Merge pull request #105 from sblunt/tau

Tau utilities and customization. Will open an issue to write documentation for tau

orbitize/basis.py

@@ -0,0 +1,65 @@
+import numpy as np
+import astropy.units as u
+
+def tau_to_t0(tau, ref_epoch, period, after_date=None):
+    """
+    Convert tau (epoch of periastron in fractional orbital period after ref epoch) to
+    T0 (date in days, usually MJD, but works with whatever system ref_epoch is given in)
+
+    Args:
+        tau (float or np.array): value of tau to convert
+        ref_epoch (float or np.array): date (in days, typically MJD) that tau is defined relative to
+        period (float or np.array): period (in years) that tau is noralized with
+        after_date (float): T0 will be the first periastron after this date. If None, use ref_epoch.
+
+    Returns:
+        t0 (float or np.array): corresponding T0 of the taus
+    """
+    period_days = period * u.year.to(u.day)
+
+    t0 = tau * (period_days) + ref_epoch
+
+    if after_date is not None:
+        num_periods = (after_date - t0)/period_days
+        num_periods = int(np.ceil(num_periods))
+
+        t0 += num_periods * period_days
+
+    return t0
+
+def t0_to_tau(t0, ref_epoch, period):
+    """
+    Convert T0 to tau
+
+    Args:
+        t0 (float or np.array): value to T0 to convert (days, typically MJD)
+        ref_epoch (float or np.array): reference epoch (in days) that tau is defined from. Same system as t0 (e.g., MJD)
+        period (float or np.array): period (in years) that tau is defined by
+
+    Returns:
+        tau (float or np.array): corresponding taus
+    """
+    tau = (t0 - ref_epoch)/(period * u.year.to(u.day))
+    tau %= 1
+
+    return tau
+
+def switch_tau_epoch(old_tau, old_epoch, new_epoch, period):
+    """
+    Convert tau to another tau that uses a different referench epoch
+
+    Args:
+        old_tau (float or np.array): old tau to convert
+        old_epoch (float or np.array): old reference epoch (days, typically MJD)
+        new_epoch (float or np.array): new reference epoch (days, same system as old_epoch)
+        period (float or np.array): orbital period (years)
+
+    Returns:
+        new_tau (float or np.array): new taus
+    """
+    period_days = period * u.year.to(u.day)
+
+    t0 = tau_to_t0(old_tau, old_epoch, period)
+    new_tau = t0_to_tau(t0, new_epoch, period)
+
+    return new_tau

orbitize/driver.py

@@ -24,14 +24,17 @@ class Driver(object):
         plx_err (float, optional): uncertainty on ``plx`` [mas]
         lnlike (str, optional): name of function in ``orbitize.lnlike`` that will
             be used to compute likelihood. (default="chi2_lnlike")
+        system_kwargs (dict, optional): ``restrict_angle_ranges``, ``ref_tau_epoch``, 
+            ``results`` for ``orbitize.system.System``. 
         mcmc_kwargs (dict, optional): ``num_temps``, ``num_walkers``, and ``num_threads``
             kwargs for ``orbitize.sampler.MCMC``
 
     Written: Sarah Blunt, 2018
     """
     def __init__(self, input_data, sampler_str,
                  num_secondary_bodies, system_mass, plx,
-                 mass_err=0, plx_err=0, lnlike='chi2_lnlike', mcmc_kwargs=None):
+                 mass_err=0, plx_err=0, lnlike='chi2_lnlike', 
+                 system_kwargs=None, mcmc_kwargs=None):
 
         # Read in data
         # Try to interpret input as a filename first
@@ -45,10 +48,13 @@ def __init__(self, input_data, sampler_str,
             except:
                 raise Exception('Invalid value of input_data for Driver')
 
+        if system_kwargs is None:
+            system_kwargs = {}
+
         # Initialize System object which stores data & sets priors
         self.system = orbitize.system.System(
             num_secondary_bodies, data_table, system_mass,
-            plx, mass_err=mass_err, plx_err=plx_err
+            plx, mass_err=mass_err, plx_err=plx_err, **system_kwargs
         )
 
         # Initialize Sampler object, which stores information about

orbitize/kepler.py

@@ -14,7 +14,7 @@
     cext = False
 
 
-def calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=None, tolerance=1e-9, max_iter=100):
+def calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=None, tau_ref_epoch=0, tolerance=1e-9, max_iter=100):
     """
     Returns the separation and radial velocity of the body given array of
     orbital parameters (size n_orbs) at given epochs (array of size n_dates)
@@ -32,6 +32,7 @@ def calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=None, tole
         plx (np.array): parallax [mas]
         mtot (np.array): total mass [Solar masses]
         mass (np.array, optional): mass of the body [Solar masses]. For planets mass ~ 0 (default)
+        tau_ref_epoch (float, optional): reference date that tau is defined with respect to (i.e., tau=0)
         tolerance (float, optional): absolute tolerance of iterative computation. Defaults to 1e-9.
         max_iter (int, optional): maximum number of iterations before switching. Defaults to 100.
 
@@ -66,7 +67,7 @@ def calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=None, tole
     mean_motion = 2*np.pi/(period) # in rad/day
 
     # # compute mean anomaly (size: n_orbs x n_dates)
-    manom = (mean_motion*epochs[:, None] - 2*np.pi*tau) % (2.0*np.pi)
+    manom = (mean_motion*(epochs[:, None] - tau_ref_epoch) - 2*np.pi*tau) % (2.0*np.pi)
 
     # compute eccentric anomalies (size: n_orbs x n_dates)
     eanom = _calc_ecc_anom(manom, ecc_arr, tolerance=tolerance, max_iter=max_iter)

orbitize/results.py

@@ -38,6 +38,7 @@ class Results(object):
             parameters in the fit (default: None).
         lnlike (np.array of float): M array of log-likelihoods corresponding to
             the orbits described in ``post`` (default: None).
+        tau_ref_epoch (float): date (in days, typically MJD) that tau is defined relative to
 
     The ``post`` array is in the following order::
 
@@ -48,14 +49,15 @@ class Results(object):
         [parallax, total mass]
 
     where 1 corresponds to the first orbiting object, 2 corresponds
-    to the second, etc. 
+    to the second, etc.
 
     Written: Henry Ngo, Sarah Blunt, 2018
     """
-    def __init__(self, sampler_name=None, post=None, lnlike=None):
+    def __init__(self, sampler_name=None, post=None, lnlike=None, tau_ref_epoch=None):
         self.sampler_name = sampler_name
         self.post = post
         self.lnlike = lnlike
+        self.tau_ref_epoch = tau_ref_epoch
 
     def add_samples(self, orbital_params, lnlikes):
         """
@@ -95,12 +97,13 @@ def save_results(self, filename, format='hdf5'):
         MCMC sampler has the ``lnlike`` attribute set. For OFTI, ``lnlike`` is None and
         it is not saved.
 
-        HDF5: ``sampler_name`` is an attribute of the root group. ``post`` and ``lnlike``
-        are datasets that are members of the root group.
+        HDF5: ``sampler_name`` and ``tau_ref_epcoh`` are attributes of the root group.
+        ``post`` and ``lnlike`` are datasets that are members of the root group.
 
         FITS: Data is saved as Binary FITS Table to the *first extension* HDU.
         After reading with something like ``hdu = astropy.io.fits.open(file)``,
         ``hdu[1].header['SAMPNAME']`` returns the ``sampler_name``.
+        ``hdu[1].header['TAU_REF']`` returns the ``tau_ref_epoch``.
         ``hdu[1].data`` returns a ``Table`` with two columns. The first column
         contains the post array, and the second column contains the lnlike array
 
@@ -110,6 +113,7 @@ def save_results(self, filename, format='hdf5'):
             hf = h5py.File(filename,'w') # Creates h5py file object
             # Add sampler_name as attribute of the root group
             hf.attrs['sampler_name']=self.sampler_name
+            hf.attrs['tau_ref_epoch'] = self.tau_ref_epoch
             # Now add post and lnlike from the results object as datasets
             hf.create_dataset('post', data=self.post)
             if self.lnlike is not None: # This property doesn't exist for OFTI
@@ -130,6 +134,7 @@ def save_results(self, filename, format='hdf5'):
                 hdu = fits.BinTableHDU.from_columns([col_post])
             # Add sampler_name to the hdu's header
             hdu.header['SAMPNAME'] = self.sampler_name
+            hdu.header['TAU_REF'] = self.tau_ref_epoch
             # Write to fits file
             hdu.writeto(filename)
         else:
@@ -156,12 +161,28 @@ def load_results(self, filename, format='hdf5', append=False):
             sampler_name = np.str(hf.attrs['sampler_name'])
             post = np.array(hf.get('post'))
             lnlike = np.array(hf.get('lnlike'))
+            # get the tau reference epoch
+            try:
+                tau_ref_epoch = float(hf.attrs['tau_ref_epoch'])
+            except KeyError:
+                # probably a old results file when reference epoch was fixed at MJD = 0
+                tau_ref_epoch = 0
+
             hf.close() # Closes file object
         elif format.lower()=='fits':
             hdu_list = fits.open(filename) # Opens file as HDUList object
             table_hdu = hdu_list[1] # Table data is in first extension
+
             # Get sampler_name from header
             sampler_name = table_hdu.header['SAMPNAME']
+
+            # get tau reference epoch
+            if 'TAU_REF' in table_hdu.header:
+                tau_ref_epoch = table_hdu.header['TAU_REF']
+            else:
+                # this is most likely an older results file where it was fixed to MJD=0
+                tau_ref_epoch = 0
+
             # Get post and lnlike arrays from column names
             post = table_hdu.data.field('post')
             # (Note: OFTI does not have lnlike so it won't be saved)
@@ -182,13 +203,22 @@ def load_results(self, filename, format='hdf5', append=False):
             # otherwise only proceed if the sampler_names match
             elif self.sampler_name != sampler_name:
                 raise Exception('Unable to append file {} to Results object. sampler_name of object and file do not match'.format(filename))
+
+            # if no tau reference epoch is set, use input file's value
+            if self.tau_ref_epoch is None:
+                self.tau_ref_epoch = tau_ref_epoch
+            # otherwise, only proceed if they are identical
+            elif self.tau_ref_epoch != tau_ref_epoch:
+                raise ValueError("Loaded data has tau reference epoch of {0} while Results object has already been initialized to {1}".format(tau_ref_epoch, self.tau_ref_epoch))
+
             # Now append post and lnlike
             self.add_samples(post,lnlike)
         else:
             # Only proceed if object is completely empty
-            if self.sampler_name is None and self.post is None and self.lnlike is None:
+            if self.sampler_name is None and self.post is None and self.lnlike is None and self.tau_ref_epoch is None:
                 self._set_sampler_name(sampler_name)
                 self.add_samples(post,lnlike)
+                self.tau_ref_epoch = tau_ref_epoch
             else:
                 raise Exception('Unable to load file {} to Results object. append is set to False but object is not empty'.format(filename))
 
@@ -281,7 +311,7 @@ def plot_corner(self, param_list=[], **corner_kwargs):
     def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
                     object_to_plot=1, start_mjd=51544.,
                     num_orbits_to_plot=100, num_epochs_to_plot=100,
-                    square_plot=True, show_colorbar=True, cmap=cmap, 
+                    square_plot=True, show_colorbar=True, cmap=cmap,
                     sep_pa_color='lightgrey', sep_pa_end_year=2025.0,
                     cbar_param='epochs'):
 
@@ -309,11 +339,11 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
             show_colorbar (Boolean): Displays colorbar to the right of the plot [True]
             cmap (matplotlib.cm.ColorMap): color map to use for making orbit tracks
                 (default: modified Purples_r)
-            sep_pa_color (string): any valid matplotlib color string, used to set the 
+            sep_pa_color (string): any valid matplotlib color string, used to set the
                 color of the orbit tracks in the Sep/PA panels (default: 'lightgrey').
             sep_pa_end_year (float): decimal year specifying when to stop plotting orbit
                 tracks in the Sep/PA panels (default: 2025.0).
-            cbar_param (string): options are the following: epochs, sma1, ecc1, inc1, aop1, 
+            cbar_param (string): options are the following: epochs, sma1, ecc1, inc1, aop1,
                 pan1, tau1. Number can be switched out. Default is epochs.
 
         Return:
@@ -336,7 +366,7 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
                 'pan': 4,
                 'tau': 5
             }
-            
+
             if cbar_param == 'epochs':
                 pass
             elif cbar_param[0:3] in dict_of_indices:
@@ -348,13 +378,13 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
                 index = dict_of_indices[cbar_param[0:3]] + 6*(object_id-1)
             else:
                 raise Exception('Invalid input; acceptable inputs include epochs, sma1, ecc1, inc1, aop1, pan1, tau1, sma2, ecc2, ...')
-            
+
 
             # Split the 2-D post array into series of 1-D arrays for each orbital parameter
             num_objects, remainder = np.divmod(self.post.shape[1],6)
             if object_to_plot > num_objects:
                 return None
-            
+
             sma = self.post[:,dict_of_indices['sma']]
             ecc = self.post[:,dict_of_indices['ecc']]
             inc = self.post[:,dict_of_indices['inc']]
@@ -399,7 +429,7 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
                 # Calculate ra/dec offsets for all epochs of this orbit
                 raoff0, deoff0, _ = kepler.calc_orbit(
                     epochs[i,:], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind],
-                    tau[orb_ind], plx[orb_ind], mtot[orb_ind], mass=mplanet[orb_ind]
+                    tau[orb_ind], plx[orb_ind], mtot[orb_ind], mass=mplanet[orb_ind], tau_ref_epoch=self.tau_ref_epoch
                 )
 
                 raoff[i,:] = raoff0
@@ -464,20 +494,20 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
             epochs_seppa = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
 
             for i in np.arange(num_orbits_to_plot):
-                
+
                 orb_ind = choose[i]
 
 
                 epochs_seppa[i,:] = np.linspace(
-                    start_mjd, 
-                    Time(sep_pa_end_year, format='decimalyear').mjd, 
+                    start_mjd,
+                    Time(sep_pa_end_year, format='decimalyear').mjd,
                     num_epochs_to_plot
                 )
 
                 # Calculate ra/dec offsets for all epochs of this orbit
                 raoff0, deoff0, _ = kepler.calc_orbit(
                     epochs_seppa[i,:], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind],
-                    tau[orb_ind], plx[orb_ind], mtot[orb_ind], mass=mplanet[orb_ind]
+                    tau[orb_ind], plx[orb_ind], mtot[orb_ind], mass=mplanet[orb_ind], tau_ref_epoch=self.tau_ref_epoch
                 )
 
                 raoff[i,:] = raoff0

orbitize/sampler.py

@@ -127,7 +127,8 @@ def __init__(self, system, like='chi2_lnlike', custom_lnlike=None):
                 self.sep_err[i], self.pa_err[i]
             )
 
-        self.epochs = np.array(self.data_table['epoch'])
+        ### this is OK, ONLY IF we are only using self.epochs for computing RA/Dec from Keplerian elements
+        self.epochs = np.array(self.data_table['epoch']) - self.system.tau_ref_epoch 
 
         # choose scale-and-rotate epoch
         self.epoch_idx = np.argmin(self.sep_err) # epoch with smallest error
@@ -136,7 +137,8 @@ def __init__(self, system, like='chi2_lnlike', custom_lnlike=None):
         self.results = orbitize.results.Results(
             sampler_name = self.__class__.__name__,
             post = None,
-            lnlike = None
+            lnlike = None,
+            tau_ref_epoch=self.system.tau_ref_epoch
         )
 
     def prepare_samples(self, num_samples):
@@ -317,7 +319,8 @@ def __init__(self, system, num_temps=20, num_walkers=1000, num_threads=1, like='
         self.results = orbitize.results.Results(
             sampler_name = self.__class__.__name__,
             post = None,
-            lnlike = None
+            lnlike = None,
+            tau_ref_epoch=system.tau_ref_epoch
         )
 
         if self.num_temps > 1:

orbitize/system.py

@@ -18,6 +18,8 @@ class System(object):
         restrict_angle_ranges (bool, optional): if True, restrict the ranges
             of the position angle of nodes and argument of periastron to [0,180)
             to get rid of symmetric double-peaks for imaging-only datasets.
+        tau_ref_epoch (float, optional): reference epoch for defining tau (MJD).
+            Default is 58849 (Jan 1, 2020).
         results (list of orbitize.results.Results): results from an orbit-fit
             will be appended to this list as a Results class
 
@@ -40,12 +42,13 @@ class System(object):
     """
     def __init__(self, num_secondary_bodies, data_table, system_mass,
                  plx, mass_err=0, plx_err=0, restrict_angle_ranges=None,
-                 results=None):
+                 tau_ref_epoch=58849, results=None):
 
         self.num_secondary_bodies = num_secondary_bodies
         self.sys_priors = []
         self.labels = []
         self.results = []
+        self.tau_ref_epoch = tau_ref_epoch
 
         #
         # Group the data in some useful ways
@@ -169,7 +172,7 @@ def compute_model(self, params_arr):
             mtot = params_arr[-1]
 
             raoff, decoff, vz = kepler.calc_orbit(
-                epochs, sma, ecc, inc, argp, lan, tau, plx, mtot
+                epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, tau_ref_epoch=self.tau_ref_epoch
             )
 
             if len(raoff[self.radec[body_num]]) > 0: # (prevent empty array dimension errors)