Merge pull request #48 from sblunt/isabel_dev

Working OFTI implementation

orbitize/driver.py

@@ -1,4 +1,6 @@
-from orbitize import read_input, system
+import orbitize.read_input
+import orbitize.system
+import orbitize.sampler
 
 
 class Driver(object):
@@ -7,7 +9,7 @@ class Driver(object):
 
     Args:
         filename (str): path to data file. See ``orbitize.read_input.py``
-        sampler (str): algorithm to use for orbit computation. 'PTMCMC' for 
+        sampler_str (str): algorithm to use for orbit computation. 'PTMCMC' for 
             Markov Chain Monte Carlo, 'OFTI' for Orbits for the Impatient.
         lnlike (str): name of function in ``orbitize.lnlike.py`` that will
             be used to compute likelihood. ["chi2_lnlike"]
@@ -20,24 +22,24 @@ class Driver(object):
 
     (written): Sarah Blunt, 2018
     """
-    def __init__(self, filename, sampler,
+    def __init__(self, filename, sampler_str,
                  num_secondary_bodies, system_mass, plx, 
                  mass_err=0, plx_err=0, lnlike='chi2_lnlike'):
 
         # Read in data
-        data_table = read_input.read_formatted_file(filename)
+        data_table = orbitize.read_input.read_formatted_file(filename)
 
         # Initialize System object which stores data & sets priors
-        self.system = system.System(
+        self.system = orbitize.system.System(
             num_secondary_bodies, data_table, system_mass, 
             plx, mass_err=mass_err, plx_err=plx_err
         )
 
         # Initialize Sampler object, which stores information about
         # the likelihood function & the algorithm used to generate
         # orbits, and has System object as an attribute.
-        sampler_func = getattr(orbitize.sampler, sampler)
-        self.sampler = sampler_func(lnlike, self.system)
+        sampler_func = getattr(orbitize.sampler, sampler_str)
+        self.sampler = sampler_func(self.system, lnlike)
 
     def compute_posteriors(self, total_orbits):
 

orbitize/kepler.py

@@ -36,9 +36,9 @@ def calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=None, tole
         max_iter (int, optional): maximum number of iterations before switching. Defaults to 100.
 
     Return:
-        raoff (np.array): array-like (n_dates x n_orbs) of RA offsets between the bodies (origin is at the other body)
-        deoff (np.array): array-like (n_dates x n_orbs) of Dec offsets between the bodies
-        vz (np.array): array-like (n_dates x n_orbs) of radial velocity offset between the bodies
+        raoff (np.array): array-like (n_dates x n_orbs) of RA offsets between the bodies (origin is at the other body) in mas
+        deoff (np.array): array-like (n_dates x n_orbs) of Dec offsets between the bodies in mas
+        vz (np.array): array-like (n_dates x n_orbs) of radial velocity offset between the bodies 
 
     Written: Jason Wang, Henry Ngo, 2018
     """

orbitize/lnlike.py

@@ -24,15 +24,30 @@ def chi2_lnlike(data, errors, model, seppa_indices):
     	function should be an array of dimension 8 x 2 x 10,000.
 
     """
+    if np.ndim(model) == 3:
+        # move M dimension to the primary axis, so that numpy knows to iterate over it
+        model = np.rollaxis(model, 2, 0) # now MxNobsx2 in dimensions
+        third_dim = True
+    elif np.ndim(model) == 2:
+        model.shape = (1,) + model.shape
+        third_dim = False
 
     chi2 = (data - model)**2/errors**2
 
     # if there are PA values, we should take the difference modulo angle wrapping
     if np.size(seppa_indices) > 0:
-        chi2[seppa_indices, 1] = np.arctan2(
-            np.sin(data[seppa_indices, 1] - model[seppa_indices, 1]), 
-            np.cos(data[seppa_indices, 1] - model[seppa_indices, 1])
+        chi2[:, seppa_indices, 1] = np.arctan2(
+            np.sin(data[seppa_indices, 1] - model[:, seppa_indices, 1]), 
+            np.cos(data[seppa_indices, 1] - model[:, seppa_indices, 1])
         )**2 / errors[seppa_indices, 1]**2
 
+    if third_dim:
+        # move M dimension back to the last axis
+        model = np.rollaxis(model, 0, 3) # now MxNobsx2 in dimensions
+        chi2 = np.rollaxis(chi2, 0, 3) # same with chi2
+    else:
+        model.shape = model.shape[1:]
+        chi2.shape = chi2.shape[1:]
+
     return chi2
 

orbitize/results.py

@@ -25,7 +25,7 @@ class Results(object):
         argument of periastron 1, position angle of nodes 1,
         epoch of periastron passage 1,
         [semimajor axis 2, eccentricity 2, etc.],
-        [total mass, parallax]
+        [parallax, total mass]
 
     where 1 corresponds to the first orbiting object, 2 corresponds
     to the second, etc. If stellar mass

orbitize/sampler.py

@@ -1,12 +1,18 @@
-import orbitize.lnlike
-import orbitize.priors
-import orbitize.results
+import numpy as np
+import astropy.units as u
+import astropy.constants as consts
 import sys
 import abc
-import numpy as np
+
 import emcee
 import ptemcee
 
+import orbitize.lnlike
+import orbitize.priors
+import orbitize.kepler
+from orbitize.system import radec2seppa
+import orbitize.results
+
 # Python 2 & 3 handle ABCs differently
 if sys.version_info[0] < 3:
     ABC = abc.ABCMeta('ABC', (), {})
@@ -17,7 +23,6 @@ class Sampler(ABC):
     """
     Abstract base class for sampler objects.
     All sampler objects should inherit from this class.
-
     (written): Sarah Blunt, 2018
     """
 
@@ -42,58 +47,186 @@ class OFTI(Sampler):
     Args:
         lnlike (string): name of likelihood function in ``lnlike.py``
         system (system.System): system.System object
+
+    (written): Isabel Angelo, Logan Pearce, Sarah Blunt 2018
     """
     def __init__(self, system, like='chi2_lnlike'):
+
         super(OFTI, self).__init__(system, like=like)
+
+        self.priors = self.system.sys_priors
+        self.tbl = self.system.data_table
+        self.radec_idx = self.system.radec[1]
+        self.seppa_idx = self.system.seppa[1]
+
+        # these are of type astropy.table.column
+        self.sep_observed = self.tbl[:]['quant1']
+        self.pa_observed = self.tbl[:]['quant2']
+        self.sep_err = self.tbl[:]['quant1_err']
+        self.pa_err = self.tbl[:]['quant2_err']
+
+        # convert RA/Dec rows to sep/PA
+        for i in self.radec_idx:
+            self.sep_observed[i], self.pa_observed[i] = radec2seppa(
+                self.sep_observed[i], self.pa_observed[i]
+            )
+            self.sep_err[i], self.pa_err[i] = radec2seppa(
+                self.sep_err[i], self.pa_err[i]
+            )
+
+        self.epochs = np.array(self.tbl['epoch'])
+
+        # choose scale-and-rotate epoch
+        self.epoch_idx = np.argmin(self.sep_err) # epoch with smallest error
+
+        # format sep/PA observations for use with the lnlike code
+        self.seppa_for_lnlike = np.column_stack((self.sep_observed, self.pa_observed))
+        self.seppa_errs_for_lnlike = np.column_stack((self.sep_err, self.pa_err))
 
     def prepare_samples(self, num_samples):
         """
         Prepare some orbits for rejection sampling. This draws random orbits
         from priors, and performs scale & rotate.
 
         Args:
-            num_samples (int): number of orbits to prepare for OFTI to run
-                rejection sampling on
+            num_samples (int): number of orbits to draw and scale & rotate for 
+            OFTI to run rejection sampling on
 
         Return:
             np.array: array of prepared samples. The first dimension has size of 
-            num_samples. This should be passed into `reject()`
+            num_samples. This should be passed into ``reject()``
         """
-        pass
-        # draw an array of num_samples smas, eccs, etc. from prior objects: prior = (some object inhertiting from priors.Prior); samples = prior.draw_samples(#)
-      #  elements = system.priors # -> this step should be done in OFTI.__init__ so it doesn't slow performance
 
-    #    for element in elements:
-     #       samples[i,j] = system.priors[element].draw_samples(num_samples)
+        # TODO: modify to work for multi-planet systems
+
+        # generate sample orbits
+        samples = np.empty([len(self.priors), num_samples])
+        for i in range(len(self.priors)): 
+            samples[i, :] = self.priors[i].draw_samples(num_samples)
 
-    def reject(self, orbit_configs):
+        # TODO: fix for the case where m_err and plx_err are nan
+        sma, ecc, inc, argp, lan, tau, plx, mtot = [s for s in samples]
+
+        period_prescale = np.sqrt(
+            4*np.pi**2*(sma*u.AU)**3/(consts.G*(mtot*u.Msun))
+        )
+        period_prescale = period_prescale.to(u.day).value
+        meananno = self.epochs[self.epoch_idx]/period_prescale - tau
+
+        # compute sep/PA of generated orbits 
+        ra, dec, vc = orbitize.kepler.calc_orbit(
+            self.epochs[self.epoch_idx], sma, ecc, inc, argp, lan, tau, plx, mtot
+        )
+        sep, pa = orbitize.system.radec2seppa(ra, dec) # sep[mas], PA[deg]  
+
+        # generate Gaussian offsets from observational uncertainties
+        sep_offset = np.random.normal(
+            0, self.sep_err[self.epoch_idx], size=num_samples
+        )
+        pa_offset =  np.random.normal(
+            0, self.pa_err[self.epoch_idx], size=num_samples
+        )
+
+        # calculate correction factors
+        sma_corr = (sep_offset + self.sep_observed[self.epoch_idx])/sep
+        lan_corr = (pa_offset + self.pa_observed[self.epoch_idx] - pa)
+
+        # perform scale-and-rotate
+        sma *= sma_corr # [AU]
+        lan += np.radians(lan_corr) # [rad] 
+        lan = lan % (2*np.pi)
+
+        period_new = np.sqrt(
+            4*np.pi**2*(sma*u.AU)**3/(consts.G*(mtot*u.Msun))
+        )
+        period_new = period_new.to(u.day).value
+
+        tau = (self.epochs[self.epoch_idx]/period_new - meananno)
+
+        # updates samples with new values of sma, pan, tau
+        samples[0,:] = sma
+        samples[4,:] = lan
+        samples[5,:] = tau
+
+        return samples
+
+
+    def reject(self, samples):
         """
-        Runs rejection sampling on some prepared samples
+        Runs rejection sampling on some prepared samples.
 
         Args:
-            orbit_configs (np.array): array of prepared samples. The first dimension has size `num_samples`. This should be the output of ``prepare_samples()``
+            samples (np.array): array of prepared samples. The first dimension 
+            has size `num_samples`. This should be the output of 
+            `prepare_samples()`.
 
         Return:
-            np.array: a subset of orbit_configs that are accepted based on the data.
-
+            np.array: a subset of `samples` that are accepted based on the 
+                data.
+            
         """
-        pass
+
+        # generate seppa for all remaining epochs
+        sma, ecc, inc, argp, lan, tau, plx, mtot = [s for s in samples]
+
+        ra, dec, vc = orbitize.kepler.calc_orbit(
+            self.epochs, sma, ecc, inc, argp, lan, tau, plx, mtot
+        )
+        sep, pa = orbitize.system.radec2seppa(ra, dec)
 
-    def run_sampler(self, total_orbits):
+        seppa_model = np.vstack(zip(sep, pa))
+        seppa_model = seppa_model.reshape((len(self.epochs), 2, len(sma)))
+
+        # compute chi2 for each orbit
+        chi2 = orbitize.lnlike.chi2_lnlike(
+            self.seppa_for_lnlike, self.seppa_errs_for_lnlike, 
+            seppa_model, self.seppa_idx
+        )
+
+        # convert to log(probability)
+        chi2_sum = np.nansum(chi2, axis=(0,1))
+        lnp = -chi2_sum/2.
+
+        # reject orbits with probability less than a uniform random number
+        random_samples = np.log(np.random.random(len(lnp)))
+        saved_orbit_idx = np.where(lnp > random_samples)[0]
+        saved_orbits = np.array([samples[:,i] for i in saved_orbit_idx])
+
+        return saved_orbits
+
+
+    def run_sampler(self, total_orbits, num_samples=10000):
         """
-        Runs OFTI until we get the number of total accepted orbits we want.
+        Runs OFTI until we get the number of total accepted orbits we want. 
 
         Args:
-            total_orbits (int): total number of accepted possible orbits that
-                are desired
+            total_orbits (int): total number of accepted orbits desired by user
+            num_samples (int): number of orbits to prepare for OFTI to run
+            rejection sampling on
 
         Return:
-            np.array: array of accepted orbits. First dimension has size ``total_orbits``.
+            output_orbits (np.array): array of accepted orbits. First dimension 
+            has size `total_orbits`.
         """
-        # this function shold first check if we have reached enough orbits, and break when we do
 
-        # put outputs of calc_orbit into format specified by mask passed from System object. Feed these arrays of data, model, and errors into lnlike.py
-        pass
+        n_orbits_saved = 0
+        output_orbits = np.empty((total_orbits, len(self.priors)))
+
+        # add orbits to `output_orbits` until `total_orbits` are saved
+        while n_orbits_saved < total_orbits:
+            samples = self.prepare_samples(num_samples)
+            accepted_orbits = self.reject(samples)
+
+            if len(accepted_orbits)==0:
+                pass
+            else:
+                n_accepted = len(accepted_orbits)
+                maxindex2save = np.min([n_accepted, total_orbits - n_orbits_saved])
+
+                output_orbits[n_orbits_saved : n_orbits_saved+n_accepted] = accepted_orbits[0:maxindex2save]
+                n_orbits_saved += maxindex2save
+
+        return np.array(output_orbits)
 
 
 class PTMCMC(Sampler):