Merge pull request #147 from sblunt/corner-in-degrees

Default corner plot now plots angles in degrees
sblunt · Nov 10, 2019 · 5f88d3a · 5f88d3a
2 parents 34856e2 + f21804b
commit 5f88d3a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 32 deletions.
diff --git a/orbitize/results.py b/orbitize/results.py
@@ -1,6 +1,7 @@
 import numpy as np
 import warnings
 import h5py
+import copy
 
 import astropy.units as u
 import astropy.constants as consts
@@ -219,30 +220,33 @@ def plot_corner(self, param_list=None, **corner_kwargs):
         default_labels = {
             'sma':'a [au]',
             'ecc':'ecc',
-            'inc':'inc [rad]',
-            'aop':'$\omega$ [rad]',
-            'pan':'$\Omega$ [rad]',
+            'inc':'inc [$^{\circ}$]',
+            'aop':'$\omega$ [$^{\circ}$]',
+            'pan':'$\Omega$ [$^{\circ}$]',
             'tau':'$\\tau$',
             'plx':'$\pi$ [mas]',
-            'mtot':'$M_T$ [Msol]',
-            'm0':'$M_0$ [Msol]',
-            'm1':'$M_1$ [Msol]',
+            'mtot':'$M_T$ [M$_{\odot}$]',
+            'm0':'$M_0$ [M$_{\odot}$]',
+            'm1':'$M_1$ [M$_{\odot}$]',
         }
 
-        if param_list is not None:
-            param_indices = []
-            for param in param_list:
-                index_num = np.where(np.array(self.labels) == param)[0][0]
-                param_indices.append(index_num)
+        if param_list is None:
+            param_list = self.labels
 
-            samples = self.post[:,param_indices] # Keep only chains for selected parameters
+        param_indices = []
+        angle_indices = []
+        for i, param in enumerate(param_list):
+            index_num = np.where(np.array(self.labels) == param)[0][0]
+            param_indices.append(index_num)
+            label_key = param_list[i]
+            if label_key.startswith('aop') or label_key.startswith('pan') or label_key.startswith('inc'):
+                angle_indices.append(index_num)
+
+        samples = copy.copy(self.post[:,param_indices]) # keep only chains for selected parameters
+        samples[:,angle_indices] = np.degrees(self.post[:,angle_indices]) # convert angles from rad to deg
 
-        else:
-            param_list = self.labels
-            samples = self.post
-            param_indices = np.arange(len(param_list))
 
-        if 'labels' not in corner_kwargs: # Use default labels if user didn't already supply them
+        if 'labels' not in corner_kwargs: # use default labels if user didn't already supply them
             reduced_labels_list = []
             for i in np.arange(len(param_indices)):
                 label_key = param_list[i]

diff --git a/tests/test_results.py b/tests/test_results.py
@@ -22,12 +22,12 @@ def simulate_orbit_sampling(n_sim_orbits):
     sma_err = 1.1  # not real
     ecc = 0.08
     ecc_err = 0.03 # not real
-    inc = 88.81
-    inc_err = 0.12
-    aop = 205.8
-    aop_err = 20.0 # not real
-    pan = 31.76
-    pan_err = 0.09
+    inc = np.radians(88.81)
+    inc_err = np.radians(0.12)
+    aop = np.radians(205.8)
+    aop_err = np.radians(20.0) # not real
+    pan = np.radians(31.76)
+    pan_err = np.radians(0.09)
     epp = 0.73
     epp_err = 0.20 # not real
     system_mass = 1.80
@@ -42,8 +42,8 @@ def simulate_orbit_sampling(n_sim_orbits):
     sim_post[:,3]=np.random.normal(aop,aop_err,n_sim_orbits)
     sim_post[:,4]=np.random.normal(pan,pan_err,n_sim_orbits)
     sim_post[:,5]=np.random.normal(epp,epp_err,n_sim_orbits)
-    sim_post[:,6]=np.random.normal(system_mass,system_mass_err,n_sim_orbits)
-    sim_post[:,7]=np.random.normal(plx,plx_err,n_sim_orbits)
+    sim_post[:,6]=np.random.normal(plx,plx_err,n_sim_orbits)
+    sim_post[:,7]=np.random.normal(system_mass,system_mass_err,n_sim_orbits)
 
     return sim_post
 
@@ -170,10 +170,10 @@ def test_plot_orbits(results_to_test):
     test_save_and_load_results(test_results, has_lnlike=False)
     test_corner_fig1, test_corner_fig2 = test_plot_corner(test_results)
     test_orbit_figs = test_plot_orbits(test_results)
-    # test_corner_fig1.savefig('test_corner1.png')
-    # test_corner_fig2.savefig('test_corner2.png')
-    # test_orbit_figs[0].savefig('test_orbit1.png')
-    # test_orbit_figs[1].savefig('test_orbit2.png')
-    # test_orbit_figs[2].savefig('test_orbit3.png')
-    # test_orbit_figs[3].savefig('test_orbit4.png')
-    # test_orbit_figs[4].savefig('test_orbit5.png')
+    test_corner_fig1.savefig('test_corner1.png')
+    test_corner_fig2.savefig('test_corner2.png')
+    test_orbit_figs[0].savefig('test_orbit1.png')
+    test_orbit_figs[1].savefig('test_orbit2.png')
+    test_orbit_figs[2].savefig('test_orbit3.png')
+    test_orbit_figs[3].savefig('test_orbit4.png')
+    test_orbit_figs[4].savefig('test_orbit5.png')