Merge pull request #86 from sblunt/expand-orbit-plots

Expand orbit plots
sblunt · Jan 3, 2019 · 6e8a5e7 · 6e8a5e7
2 parents 5b7d268 + e70fdf5
commit 6e8a5e7
Show file tree

Hide file tree

Showing 9 changed files with 363 additions and 144 deletions.
diff --git a/docs/tutorials/MCMC_tutorial.ipynb b/docs/tutorials/MCMC_tutorial.ipynb
diff --git a/docs/tutorials/MCMC_vs_OFTI.ipynb b/docs/tutorials/MCMC_vs_OFTI.ipynb
diff --git a/docs/tutorials/Modifying_Priors.ipynb b/docs/tutorials/Modifying_Priors.ipynb
diff --git a/docs/tutorials/OFTI_tutorial.ipynb b/docs/tutorials/OFTI_tutorial.ipynb
diff --git a/docs/tutorials/Plotting_tutorial.ipynb b/docs/tutorials/Plotting_tutorial.ipynb
diff --git a/orbitize/priors.py b/orbitize/priors.py
@@ -100,7 +100,7 @@ def compute_lnprob(self, element_array):
         lnprob = -0.5*np.log(2.*np.pi*self.sigma) - 0.5*((element_array - self.mu) / self.sigma)**2
 
         if self.no_negatives:
-
+            
             bad_samples = np.where(element_array < 0)[0]
             lnprob[bad_samples] = -np.inf
 

diff --git a/orbitize/results.py b/orbitize/results.py
@@ -7,6 +7,7 @@
 import matplotlib.colors as colors
 import corner
 import orbitize.kepler as kepler
+import orbitize.system
 import h5py
 from astropy.io import fits
 from astropy.time import Time
@@ -40,7 +41,7 @@ class Results(object):
         [parallax, total mass]
 
     where 1 corresponds to the first orbiting object, 2 corresponds
-    to the second, etc. If stellar mass
+    to the second, etc. 
 
     Written: Henry Ngo, Sarah Blunt, 2018
     """
@@ -269,10 +270,12 @@ def plot_corner(self, param_list=[], **corner_kwargs):
         figure = corner.corner(samples, **corner_kwargs)
         return figure
 
+
     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):
         """
         Plots one orbital period for a select number of fitted orbits
         for a given object, with line segments colored according to time
@@ -297,9 +300,15 @@ 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 
+                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).
+
 
         Return:
-            ``matplotlib.pyplot.Figure``: the orbit plot if input is valid, None otherwise
+            ``matplotlib.pyplot.Figure``: the orbit plot if input is valid, ``None`` otherwise
+
 
         (written): Henry Ngo, Sarah Blunt, 2018
         """
@@ -349,6 +358,7 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
             period = period.to(u.day).value
             # Create an epochs array to plot num_epochs_to_plot points over one orbital period
             epochs[i,:] = np.linspace(start_mjd, float(start_mjd+period[orb_ind]), num_epochs_to_plot)
+
             # 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],
@@ -367,7 +377,9 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
         )
 
         # Create figure for orbit plots
-        fig, ax = plt.subplots()
+        fig = plt.figure(figsize=(14,6))
+
+        ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6)
 
         # Plot each orbit (each segment between two points coloured using colormap)
         for i in np.arange(num_orbits_to_plot):
@@ -379,7 +391,7 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
             lc.set_array(epochs[i,:])
             ax.add_collection(lc)
 
-        # Modify the axes
+        # modify the axes
         if square_plot:
             adjustable_param='datalim'
         else:
@@ -390,14 +402,54 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
         ax.locator_params(axis='x', nbins=6)
         ax.locator_params(axis='y', nbins=6)
 
-        # Add colorbar
+        # add colorbar
         if show_colorbar:
-            sm = mpl.cm.ScalarMappable(cmap=cmap, norm=norm_yr)
-            sm.set_array([]) # magic? (just needs to *not* be None)
-            cbar = fig.colorbar(sm, format='%g')
-        # Alternative implementation example for right-hand colorbar
-        # fig.subplots_adjust(right=0.8)
-        # cbar_ax = fig.add_axes([0.825, 0.15, 0.05, 0.7]) # xpos, ypos, width, height, in fraction of figure size
-        # cbar = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical')
+            cbar_ax = fig.add_axes([0.47, 0.15, 0.015, 0.7]) # xpos, ypos, width, height, in fraction of figure size
+            cbar = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical')
+
+        # plot sep/PA zoom-in panels
+        ax1 = plt.subplot2grid((2, 14), (0, 9), colspan=6)
+        ax2 = plt.subplot2grid((2, 14), (1, 9), colspan=6)
+        ax2.set_ylabel('PA [$^{{\\circ}}$]')
+        ax1.set_ylabel('$\\rho$ [mas]')
+        ax2.set_xlabel('Epoch')
+
+        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, 
+                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]
+            )
+
+            raoff[i,:] = raoff0
+            deoff[i,:] = deoff0
+
+            yr_epochs = Time(epochs_seppa[i,:],format='mjd').decimalyear
+            plot_epochs = np.where(yr_epochs <= sep_pa_end_year)[0]
+            yr_epochs = yr_epochs[plot_epochs]
+
+            seps, pas = orbitize.system.radec2seppa(raoff[i,:], deoff[i,:])
+
+            plt.sca(ax1)
+            plt.plot(yr_epochs, seps, color=sep_pa_color)
+
+            plt.sca(ax2)
+            plt.plot(yr_epochs, pas, color=sep_pa_color)
+
+        ax1.locator_params(axis='x', nbins=6)
+        ax1.locator_params(axis='y', nbins=6)
+        ax2.locator_params(axis='x', nbins=6)
+        ax2.locator_params(axis='y', nbins=6)
 
         return fig
diff --git a/orbitize/sampler.py b/orbitize/sampler.py
@@ -31,7 +31,7 @@ class Sampler(ABC):
     def __init__(self, system, like='chi2_lnlike'):
         self.system = system
 
-        # check if likliehood fuction is a string of a function
+        # check if `like` is a string or a function
         if callable(like):
             self.lnlike = like
         else:

diff --git a/tests/test_priors.py b/tests/test_priors.py
@@ -74,3 +74,4 @@ def test_compute_lnprob():
 	test_draw_samples()
 
 
+