Merge pull request #239 from sblunt/plot_radec_pts_correctly

fixed a bug that was causing RA/Dec data points to display wrong in...

orbitize/example_data/test_val_radec.csv

@@ -1,5 +1,6 @@
-epoch,object,raoff,raoff_err,decoff,decoff_err
-1234,1,0.010,0.005,0.50,0.05
-1235,1,0.015,0.005,0.60,0.05
-1236,1,0.020,0.005,0.70,0.05
-1237,0,0.025,0.005,0.80,0.05
+epoch,object,raoff,raoff_err,decoff,decoff_err,sep,sep_err,pa,pa_err
+1234,1,0.010,0.005,0.50,0.05,,,,
+1235,1,0.015,0.005,0.60,0.05,,,,
+1236,1,0.020,0.005,0.70,0.05,,,,
+1237,0,0.025,0.005,0.80,0.05,,,,
+1238,0,,,,,0.025,0.005,0.80,0.05

orbitize/results.py

@@ -526,10 +526,37 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
             data=self.data
             astr_inds=np.where((~np.isnan(data['quant1'])) & (~np.isnan(data['quant2'])))
             astr_epochs=data['epoch'][astr_inds]
-            sep_data,sep_err=data['quant1'][astr_inds],data['quant1_err'][astr_inds]
-            pa_data,pa_err=data['quant2'][astr_inds],data['quant2_err'][astr_inds]
 
-
+            radec_inds = np.where(data['quant_type'] == 'radec')
+            seppa_inds = np.where(data['quant_type'] == 'seppa')
+
+            sep_data, sep_err=data['quant1'][seppa_inds],data['quant1_err'][seppa_inds]
+            pa_data, pa_err=data['quant2'][seppa_inds],data['quant2_err'][seppa_inds]
+
+            if len(radec_inds[0] > 0):
+
+
+                sep_from_ra_data, pa_from_dec_data = orbitize.system.radec2seppa(
+                    data['quant1'][radec_inds], data['quant2'][radec_inds]
+                )
+
+                num_radec_pts = len(radec_inds[0])
+                sep_err_from_ra_data = np.empty(num_radec_pts)
+                pa_err_from_dec_data = np.empty(num_radec_pts)
+                for j in np.arange(num_radec_pts):
+
+                    sep_err_from_ra_data[j], pa_err_from_dec_data[j], _ = orbitize.system.transform_errors(
+                        np.array(data['quant1'][radec_inds][j]), np.array(data['quant2'][radec_inds][j]), 
+                        np.array(data['quant1_err'][radec_inds][j]), np.array(data['quant2_err'][radec_inds][j]), 
+                        np.array(data['quant12_corr'][radec_inds][j]), orbitize.system.radec2seppa
+                    )
+
+                sep_data = np.append(sep_data, sep_from_ra_data)
+                sep_err = np.append(sep_err, sep_err_from_ra_data)
+
+                pa_data = np.append(pa_data, pa_from_dec_data)
+                pa_err = np.append(pa_err, pa_err_from_dec_data)
+
             # Plot each orbit (each segment between two points coloured using colormap)
             for i in np.arange(num_orbits_to_plot):
                 points = np.array([raoff[i, :], deoff[i, :]]).T.reshape(-1, 1, 2)

tests/test_results.py

@@ -52,13 +52,17 @@ def simulate_orbit_sampling(n_sim_orbits):
     return sim_post
 
 
-def test_init_and_add_samples():
+def test_init_and_add_samples(radec_input=False):
     """
     Tests object creation and add_samples() with some simulated posterior samples
     Returns results.Results object
     """
 
-    input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv')
+    if radec_input:
+        input_file = os.path.join(orbitize.DATADIR, 'test_val_radec.csv')
+    else:
+        input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv')
+
     data = orbitize.read_input.read_file(input_file)
 
     # Create object
@@ -194,16 +198,17 @@ def test_plot_orbits(results_to_test):
     assert Figure5 is not None
     return (Figure1, Figure2, Figure3, Figure4, Figure5)
 
-
 if __name__ == "__main__":
     test_results = test_init_and_add_samples()
+    test_results_radec = test_init_and_add_samples(radec_input=True)
 
     test_save_and_load_results(test_results, has_lnlike=True)
     test_save_and_load_results(test_results, has_lnlike=True)
     test_save_and_load_results(test_results, has_lnlike=False)
     test_save_and_load_results(test_results, has_lnlike=False)
     test_corner_fig1, test_corner_fig2, test_corner_fig3 = test_plot_corner(test_results)
     test_orbit_figs = test_plot_orbits(test_results)
+    test_orbit_figs = test_plot_orbits(test_results_radec)
     test_corner_fig1.savefig('test_corner1.png')
     test_corner_fig2.savefig('test_corner2.png')
     test_corner_fig3.savefig('test_corner3.png')