Merge 0e724db into 997647c

orbitize/sampler.py

@@ -88,11 +88,11 @@ class OFTI(Sampler):
     Args:
         like (string): name of likelihood function in ``lnlike.py``
         system (system.System): ``system.System`` object
-        custom_lnlike (func): ability to include an addition custom likelihood function in the fit. 
+        custom_lnlike (func): ability to include an addition custom likelihood function in the fit.
             the function looks like ``clnlikes = custon_lnlike(params)`` where ``params is a RxM array
-            of fitting parameters, where R is the number of orbital paramters (can be passed in system.compute_model()), 
+            of fitting parameters, where R is the number of orbital paramters (can be passed in system.compute_model()),
             and M is the number of orbits we need model predictions for. It returns ``clnlikes`` which is an array of
-            length M, or it can be a single float if M = 1. 
+            length M, or it can be a single float if M = 1.
 
     Written: Isabel Angelo, Sarah Blunt, Logan Pearce, 2018
     """
@@ -102,33 +102,21 @@ def __init__(self, system, like='chi2_lnlike', custom_lnlike=None):
 
         # compute priors and columns containing ra/dec and sep/pa
         self.priors = self.system.sys_priors
-        self.radec_idx = self.system.radec[1]
-        self.seppa_idx = self.system.seppa[1]
-
-        # store input table and table with values used by OFTI
-        self.input_table = self.system.data_table
-        self.data_table = self.system.data_table.copy()
-
-        # these are of type astropy.table.column
-        self.sep_observed = self.data_table[:]['quant1'].copy()
-        self.pa_observed = self.data_table[:]['quant2'].copy()
-        self.sep_err = self.data_table[:]['quant1_err'].copy()
-        self.pa_err = self.data_table[:]['quant2_err'].copy()
 
         # convert RA/Dec rows to sep/PA
-        if len(self.radec_idx) > 0:
+        body_num = 1 # the first planet; MODIFY THIS LATER FOR MULTIPLE PLANETS
+        if len(self.system.radec[body_num]) > 0:
             print('Converting ra/dec data points in data_table to sep/pa. Original data are stored in input_table.')
-
-        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.system.convert_data_table_radec2seppa(body_num=body_num)
+
+        # these are of type astropy.table.column
+        self.sep_observed = self.system.data_table[:]['quant1'].copy()
+        self.pa_observed = self.system.data_table[:]['quant2'].copy()
+        self.sep_err = self.system.data_table[:]['quant1_err'].copy()
+        self.pa_err = self.system.data_table[:]['quant2_err'].copy()
 
         ### 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 
+        self.epochs = np.array(self.system.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
@@ -288,7 +276,7 @@ class MCMC(Sampler):
     """
     MCMC sampler. Supports either parallel tempering or just regular MCMC. Parallel tempering will be run if ``num_temps`` > 1
     Parallel-Tempered MCMC Sampler uses ptemcee, a fork of the emcee Affine-infariant sampler
-    Affine-Invariant Ensemble MCMC Sampler uses emcee. 
+    Affine-Invariant Ensemble MCMC Sampler uses emcee.
 
     .. Warning:: may not work well for multi-modal distributions
 
@@ -299,11 +287,11 @@ class MCMC(Sampler):
         num_walkers (int): number of walkers at each temperature (default=1000)
         num_threads (int): number of threads to use for parallelization (default=1)
         like (str): name of likelihood function in ``lnlike.py``
-        custom_lnlike (func): ability to include an addition custom likelihood function in the fit. 
+        custom_lnlike (func): ability to include an addition custom likelihood function in the fit.
             the function looks like ``clnlikes = custon_lnlike(params)`` where ``params is a RxM array
-            of fitting parameters, where R is the number of orbital paramters (can be passed in system.compute_model()), 
+            of fitting parameters, where R is the number of orbital paramters (can be passed in system.compute_model()),
             and M is the number of orbits we need model predictions for. It returns ``clnlikes`` which is an array of
-            length M, or it can be a single float if M = 1. 
+            length M, or it can be a single float if M = 1.
 
     Written: Jason Wang, Henry Ngo, 2018
     """
@@ -464,9 +452,9 @@ def run_sampler(self, total_orbits, burn_steps=0, thin=1):
         except UnboundLocalError: # 0 step burn-in (pos is not defined)
             pass
         print('Burn in complete')
-        
+
         nsteps = int(np.ceil(total_orbits / self.num_walkers))
-        
+
         assert (nsteps > 0), 'Total_orbits must be greater than num_walkers.'
 
         i=0
@@ -478,7 +466,7 @@ def run_sampler(self, total_orbits, burn_steps=0, thin=1):
         print('')
 
         self.curr_pos = pos
-        
+
         # TODO: Need something here to pick out temperatures, just using lowest one for now
         self.chain = sampler.chain
 
@@ -492,9 +480,9 @@ def run_sampler(self, total_orbits, burn_steps=0, thin=1):
 
         # include fixed parameters in posterior
         self.post = self._fill_in_fixed_params(self.post)
-        
+
         self.results.add_samples(self.post,self.lnlikes)
 
         print('Run complete')
-        
+
         return sampler

orbitize/system.py

@@ -55,6 +55,8 @@ def __init__(self, num_secondary_bodies, data_table, system_mass,
         #
 
         self.data_table = data_table
+        # Creates a copy of the input in case data_table needs to be modified
+        self.input_table = self.data_table.copy()
 
         # List of arrays of indices corresponding to each body
         self.body_indices = []
@@ -190,6 +192,34 @@ def compute_model(self, params_arr):
 
         return model
 
+    def convert_data_table_radec2seppa(self,body_num=1):
+        """
+        Converts rows of self.data_table given in radec to seppa.
+        Note that self.input_table remains unchanged.
+
+        Args:
+            body_num (int): which object to convert (1 = first planet)
+        """
+        for i in self.radec[body_num]: # Loop through rows where input provided in radec
+            # Get ra/dec values
+            ra = self.data_table['quant1'][i]
+            ra_err = self.data_table['quant1_err'][i]
+            dec = self.data_table['quant2'][i]
+            dec_err = self.data_table['quant2_err'][i]
+            # Convert to sep/PA
+            sep, pa = radec2seppa(ra,dec)
+            sep_err, pa_err = radec2seppa(ra_err,dec_err)
+            # Update data_table
+            self.data_table['quant1'][i]=sep
+            self.data_table['quant1_err'][i]=sep_err
+            self.data_table['quant2'][i]=pa
+            self.data_table['quant2_err'][i]=pa_err
+            self.data_table['quant_type'][i]='seppa'
+            # Update self.radec and self.seppa arrays
+            self.radec[body_num]=np.delete(self.radec[body_num],np.where(self.radec[body_num]==i)[0])
+            self.seppa[body_num]=np.append(self.seppa[body_num],i)
+
+
     def add_results(self, results):
         """
         Adds an orbitize.results.Results object to the list in system.results

tests/test_OFTI.py

@@ -15,22 +15,22 @@
 input_file_1epoch = os.path.join(testdir, 'GJ504_1epoch.csv')
 
 def test_scale_and_rotate():
-    
+
     # perform scale-and-rotate
     myDriver = orbitize.driver.Driver(input_file, 'OFTI',
     1, 1.22, 56.95,mass_err=0.08, plx_err=0.26)
-    
+
     s = myDriver.sampler
     samples = s.prepare_samples(100)
 
     sma,ecc,inc,argp,lan,tau,plx,mtot = [samp for samp in samples]
-    
+
     ra, dec, vc = orbitize.kepler.calc_orbit(s.epochs, sma, ecc, inc, argp, lan, tau, plx, mtot)
     sep, pa = orbitize.system.radec2seppa(ra, dec)
     sep_sar, pa_sar = np.median(sep[s.epoch_idx]), np.median(pa[s.epoch_idx])
-    
+
     # test to make sure sep and pa scaled to scale-and-rotate epoch
-    sar_epoch = s.data_table[s.epoch_idx]
+    sar_epoch = s.system.data_table[s.epoch_idx]
     assert sep_sar == pytest.approx(sar_epoch['quant1'], abs=sar_epoch['quant1_err'])
     assert pa_sar == pytest.approx(sar_epoch['quant2'], abs=sar_epoch['quant2_err'])
 
@@ -39,7 +39,7 @@ def test_scale_and_rotate():
 
     # test orbit plot generation
     s.results.plot_orbits(start_mjd=s.epochs[0])
-    
+
     samples = s.results.post
     sma = samples[:,0]
     ecc = samples[:,1]
@@ -53,25 +53,25 @@ def test_scale_and_rotate():
     ra, dec, vc = orbitize.kepler.calc_orbit(s.epochs, sma, ecc, inc, argp, lan, tau, plx, mtot)
     sep, pa = orbitize.system.radec2seppa(ra, dec)
     sep_sar, pa_sar = np.median(sep[s.epoch_idx]), np.median(pa[s.epoch_idx])
-    
+
     # test to make sure sep and pa scaled to scale-and-rotate epoch
     assert sep_sar == pytest.approx(sar_epoch['quant1'], abs=sar_epoch['quant1_err'])
     assert pa_sar == pytest.approx(sar_epoch['quant2'], abs=sar_epoch['quant2_err'])
 
 
 def test_run_sampler():
-    
+
     # initialize sampler
     myDriver = orbitize.driver.Driver(input_file, 'OFTI',
     1, 1.22, 56.95,mass_err=0.08, plx_err=0.26)
-    
+
     s = myDriver.sampler
 
     # change eccentricity prior
     myDriver.system.sys_priors[1] = priors.LinearPrior(-2.18, 2.01)
-    
+
     # test num_samples=1
-    s.run_sampler(0,num_samples=1)    
+    s.run_sampler(0,num_samples=1)
 
     # test to make sure outputs are reasonable
     orbits = s.run_sampler(1000)
@@ -81,29 +81,29 @@ def test_run_sampler():
     sma = np.median([x[idx['sma1']] for x in orbits])
     ecc = np.median([x[idx['ecc1']] for x in orbits])
     inc = np.median([x[idx['inc1']] for x in orbits])
-    
+
     # expected values from Blunt et al. (2017)
     sma_exp = 48.
     ecc_exp = 0.19
     inc_exp = np.radians(140)
-    
+
     # test to make sure OFTI values are within 20% of expectations
     assert sma == pytest.approx(sma_exp, abs=0.2*sma_exp)
     assert ecc == pytest.approx(ecc_exp, abs=0.2*ecc_exp)
     assert inc == pytest.approx(inc_exp, abs=0.2*inc_exp)
-        
+
     # test with only one epoch
     myDriver = orbitize.driver.Driver(input_file_1epoch, 'OFTI',
     1, 1.22, 56.95,mass_err=0.08, plx_err=0.26)
     s = myDriver.sampler
     s.run_sampler(1)
     print()
-    
+
 def test_fixed_sys_params_sampling():
     # test in case of fixed mass and parallax
     myDriver = orbitize.driver.Driver(input_file, 'OFTI',
     1, 1.22, 56.95)
-    
+
     s = myDriver.sampler
     samples = s.prepare_samples(100)
     assert np.all(samples[-1] == s.priors[-1])
@@ -113,4 +113,4 @@ def test_fixed_sys_params_sampling():
 if __name__ == "__main__":
     test_scale_and_rotate()
     test_run_sampler()
-    print("Done!")
+    print("Done!")

tests/test_system.py

@@ -36,5 +36,22 @@ def test_add_and_clear_results():
     test_system.add_results(test_results)
     assert len(test_system.results)==1
 
+def test_convert_data_table_radec2seppa():
+    num_secondary_bodies=1
+    testdir = os.path.dirname(os.path.abspath(__file__))
+    input_file = os.path.join(testdir, 'test_val.csv')
+    data_table=read_input.read_file(input_file)
+    system_mass=1.0
+    plx=10.0
+    mass_err=0.1
+    plx_err=1.0
+    # Initialize System object
+    test_system = system.System(
+        num_secondary_bodies, data_table, system_mass,
+        plx, mass_err=mass_err, plx_err=plx_err
+    )
+    test_system.convert_data_table_radec2seppa()
+
 if __name__ == '__main__':
     test_add_and_clear_results()
+    test_convert_data_table_radec2seppa()