Merge 8f0f17f into 3eefbce

orbitize/plot.py

@@ -669,4 +669,287 @@ def plot_orbits(results, object_to_plot=1, start_mjd=51544.,
         ax2.locator_params(axis='x', nbins=6)
         ax2.locator_params(axis='y', nbins=6)
 
-    return fig
+    return fig
+
+def residual(myResults, object_to_plot=1, start_mjd=2000.0,
+                num_orbits_to_plot=100, num_epochs_to_plot=100, sep_pa_color='lightgrey',
+                sep_pa_end_year=2025.0, cbar_param='Epoch [year]',
+                mod180=False):
+
+    """
+    Plots residuals for a set of orbits
+
+    Args:
+        myResults: dataset
+        object_to_plot (int): which object to plot (default: 1)
+        start_mjd (float): MJD in which to start plotting orbits (default: 51544,
+            the year 2000)
+        num_orbits_to_plot (int): number of orbits to plot (default: 100)
+        num_epochs_to_plot (int): number of points to plot per orbit (default: 100)
+        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).
+        cbar_param (string): options are the following: 'Epoch [year]', 'sma1', 'ecc1', 'inc1', 'aop1',
+            'pan1', 'tau1', 'plx. Number can be switched out. Default is Epoch [year].
+        mod180 (Bool): if True, PA will be plotted in range [180, 540]. Useful for plotting short
+            arcs with PAs that cross 360 deg during observations (default: False)
+
+    Return:
+        ``matplotlib.pyplot.Figure``: the residual plots
+
+    """
+    data = myResults.data[myResults.data['object'] == object_to_plot]
+
+    possible_cbar_params = [
+                'sma',
+                'ecc',
+                'inc',
+                'aop'
+                'pan',
+                'tau',
+                'plx'
+            ]
+    num_orbits = len(myResults.post[:, 0])
+    if num_orbits_to_plot > num_orbits:
+        num_orbits_to_plot = num_orbits
+    choose = np.random.randint(0, high=num_orbits, size=num_orbits_to_plot)
+
+    standard_post = []
+    if myResults.sampler_name == 'MCMC':
+                # Convert the randomly chosen posteriors to standard keplerian set
+        for i in np.arange(num_orbits_to_plot):
+            orb_ind = choose[i]
+            param_set = np.copy(myResults.post[orb_ind])
+            standard_post.append(myResults.basis.to_standard_basis(param_set))
+    else: # For OFTI, posteriors are already converted
+        for i in np.arange(num_orbits_to_plot):
+            orb_ind = choose[i]
+            standard_post.append(myResults.post[orb_ind])
+
+    standard_post = np.array(standard_post)
+
+    sma = standard_post[:, myResults.standard_param_idx['sma{}'.format(object_to_plot)]]
+    ecc = standard_post[:, myResults.standard_param_idx['ecc{}'.format(object_to_plot)]]
+    inc = standard_post[:, myResults.standard_param_idx['inc{}'.format(object_to_plot)]]
+    aop = standard_post[:, myResults.standard_param_idx['aop{}'.format(object_to_plot)]]
+    pan = standard_post[:, myResults.standard_param_idx['pan{}'.format(object_to_plot)]]
+    tau = standard_post[:, myResults.standard_param_idx['tau{}'.format(object_to_plot)]]
+    plx = standard_post[:, myResults.standard_param_idx['plx']]
+
+    if 'mtot' in myResults.labels:
+        mtot = standard_post[:, myResults.standard_param_idx['mtot']]
+    elif 'm0' in myResults.labels:
+        m0 = standard_post[:, myResults.standard_param_idx['m0']]
+        m1 = standard_post[:, myResults.standard_param_idx['m{}'.format(object_to_plot)]]
+        mtot = m0 + m1
+
+    raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
+    deoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
+    vz_star = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
+    epochs = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
+
+    for i in np.arange(num_orbits_to_plot):
+        # Compute period (from Kepler's third law)
+        period = np.sqrt(4*np.pi**2.0*(sma*u.AU)**3/(consts.G*(mtot*u.Msun)))
+        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(Time(start_mjd,format='decimalyear').mjd, float(
+            Time(start_mjd,format='decimalyear').mjd+period[i]), num_epochs_to_plot)
+
+        # Calculate ra/dec offsets for all epochs of this orbit
+        raoff0, deoff0, _ = kepler.calc_orbit(
+            epochs[i, :], sma[i], ecc[i], inc[i], aop[i], pan[i],
+            tau[i], plx[i], mtot[i], tau_ref_epoch=myResults.tau_ref_epoch
+        )
+
+        raoff[i, :] = raoff0
+        deoff[i, :] = deoff0
+
+    astr_inds=np.where((~np.isnan(data['quant1'])) & (~np.isnan(data['quant2'])))        
+    astr_epochs=data['epoch'][astr_inds] 
+
+    radec_inds = np.where(data['quant_type'] == 'radec')
+    seppa_inds = np.where(data['quant_type'] == 'seppa')
+
+    # transform RA/Dec points to Sep/PA
+    sep_data = np.copy(data['quant1'])  
+    sep_err = np.copy(data['quant1_err'])
+    pa_data = np.copy(data['quant2'])
+    pa_err = np.copy(data['quant2_err'])
+
+    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[radec_inds] = sep_from_ra_data
+        sep_err[radec_inds] = sep_err_from_ra_data
+
+        pa_data[radec_inds] = pa_from_dec_data
+        pa_err[radec_inds] = pa_err_from_dec_data
+
+    # Transform Sep/PA points to RA/Dec
+    ra_data = np.copy(data['quant1'])
+    ra_err = np.copy(data['quant1_err'])
+    dec_data = np.copy(data['quant2'])
+    dec_err = np.copy(data['quant2_err'])
+
+    if len(seppa_inds[0] > 0):
+
+        ra_from_seppa_data, dec_from_seppa_data = orbitize.system.seppa2radec(
+            data['quant1'][seppa_inds], data['quant2'][seppa_inds]
+        )
+
+        num_seppa_pts = len(seppa_inds[0])
+        ra_err_from_seppa_data = np.empty(num_seppa_pts)
+        dec_err_from_seppa_data = np.empty(num_seppa_pts)
+        for j in np.arange(num_seppa_pts):
+
+            ra_err_from_seppa_data[j], dec_err_from_seppa_data[j], _ = orbitize.system.transform_errors(
+                np.array(data['quant1'][seppa_inds][j]), np.array(data['quant2'][seppa_inds][j]), 
+                np.array(data['quant1_err'][seppa_inds][j]), np.array(data['quant2_err'][seppa_inds][j]), 
+                np.array(data['quant12_corr'][seppa_inds][j]), orbitize.system.seppa2radec
+            )
+
+        ra_data[seppa_inds] = ra_from_seppa_data
+        ra_err[seppa_inds] = ra_err_from_seppa_data
+
+        dec_data[seppa_inds] = dec_from_seppa_data
+        dec_err[seppa_inds] = dec_err_from_seppa_data
+
+        epochs_seppa = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
+
+    raoff = []
+    deoff = []
+    seps = []
+    pas = []
+    raoff_100 = []
+    deoff_100 = []
+    seps_100 = []
+    pas_100 = []
+    for i in np.arange(num_orbits_to_plot):
+
+        epochs_seppa[i, :] = np.linspace(
+            Time(start_mjd, format='decimalyear').mjd ,
+            Time(sep_pa_end_year, format='decimalyear').mjd,
+            num_epochs_to_plot
+        )
+
+        raoff0, deoff0, _ = kepler.calc_orbit(
+            astr_epochs, sma[i], ecc[i], inc[i], aop[i], pan[i],
+            tau[i], plx[i], mtot[i], tau_ref_epoch=myResults.tau_ref_epoch
+        )
+
+        raoff2, deoff2, _ = kepler.calc_orbit(
+            epochs_seppa[0], sma[i], ecc[i], inc[i], aop[i], pan[i],
+            tau[i], plx[i], mtot[i], tau_ref_epoch=myResults.tau_ref_epoch
+        )
+
+        raoff.append(raoff0)
+        deoff.append(deoff0)
+        raoff_100.append(raoff2)
+        deoff_100.append(deoff2)
+
+
+        seps1 = []
+        pas1 = []
+
+        for j in range(len(astr_epochs)):
+
+            seps0, pas0 = orbitize.system.radec2seppa(raoff[i][j], deoff[i][j], mod180=mod180)
+
+            seps1.append(seps0)
+            pas1.append(pas0)
+
+        seps.append(seps1)
+        pas.append(pas1)
+
+        seps2 = []
+        pas2 = []
+
+        for j in range(len(epochs_seppa[0])):
+
+            seps0_100, pas0_100 = orbitize.system.radec2seppa(raoff_100[i][j], deoff_100[i][j], mod180=mod180)
+
+            seps2.append(seps0_100)
+            pas2.append(pas0_100)
+
+        seps_100.append(seps2)
+        pas_100.append(pas2)
+
+    yr_epochs = Time(astr_epochs, format='mjd').decimalyear
+    yr_epochs2 = Time(epochs_seppa[i, :], format='mjd').decimalyear
+
+    seps = np.array(seps)
+    pas = np.array(pas)
+    seps_100 = np.array(seps_100)
+    pas_100 = np.array(pas_100)
+
+    median_seps = []
+    median_pas = []
+    median_seps_100 = []
+    median_pas_100 = []
+
+    seps_T = np.transpose(seps)
+    pas_T = np.transpose(pas)
+    seps_100_T = np.transpose(seps_100)
+    pas_100_T = np.transpose(pas_100)
+
+    for j in range(len(epochs_seppa[0])):
+        median_seps2 = np.median(seps_100_T[j])
+        median_pas2 = np.median(pas_100_T[j])
+
+        median_seps_100.append(median_seps2)
+        median_pas_100.append(median_pas2)
+
+
+    for j in range(len(astr_epochs)):
+        median_seps1 = np.median(seps_T[j])
+        median_pas1 = np.median(pas_T[j])
+
+        median_seps.append(median_seps1)
+        median_pas.append(median_pas1)
+
+    orbits_to_plot = np.linspace(0, num_orbits_to_plot-1, 100)
+
+    residual_seps = median_seps - sep_data
+    residual_pas = median_pas - pa_data
+
+    fig, axes = plt.subplots(1, 2, figsize=(20, 10))
+
+
+    axes[0].errorbar(yr_epochs, residual_seps, yerr = sep_err, xerr = None, fmt = 'o', ms = 5,
+                    linestyle='',c='purple',zorder=10, capsize=2)
+    for i in range(len(orbits_to_plot)):
+        residual_seps_100 = median_seps_100 - seps_100[int(orbits_to_plot[i])]
+        axes[0].plot(yr_epochs2, residual_seps_100, color=sep_pa_color, zorder=1)
+    axes[0].axhline(y = 0, color = 'black', linestyle = '-')
+    axes[0].set_ylabel('Residual $\\rho$ [mas]', fontsize=18)
+    axes[0].set_xlabel('Epoch', fontsize=18)
+    axes[0].xaxis.set_tick_params(labelsize = 15)
+    axes[0].yaxis.set_tick_params(labelsize = 15)
+
+    axes[1].errorbar(yr_epochs, residual_pas, yerr = pa_err, xerr = None, fmt = 'o', ms = 5,
+                     linestyle='',c='purple',zorder=10, capsize=2)
+    for i in range(len(orbits_to_plot)):
+        residual_pas_100 = median_pas_100 - pas_100[int(orbits_to_plot[i])]
+        axes[1].plot(yr_epochs2, residual_pas_100, color=sep_pa_color, zorder=1)
+    axes[1].axhline(y = 0, color = 'black', linestyle = '-')
+    axes[1].set_ylabel('Residual PA [$^{{\\circ}}$]', fontsize=18)
+    axes[1].set_xlabel('Epoch', fontsize=18)
+    axes[1].xaxis.set_tick_params(labelsize = 15)
+    axes[1].yaxis.set_tick_params(labelsize = 15)