Skip to content

Commit

Permalink
Merge pull request #260 from sblunt/rv_plot_fix
Browse files Browse the repository at this point in the history
RV Plotting Fix
  • Loading branch information
semaphoreP committed Sep 1, 2021
2 parents d8ce41a + 8f4845c commit 44f86a5
Show file tree
Hide file tree
Showing 3 changed files with 342 additions and 23 deletions.
169 changes: 169 additions & 0 deletions orbitize/example_data/HR7672_joint.csv
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
epoch,object,rv,rv_err,sep,sep_err,pa,pa_err,instrument
47047.21939999983,0,0.27071,0.0008100000000000001,,,,,L
47373.320999999996,0,0.23288,0.00038,,,,,L
47373.4402999999,0,0.22246000000000002,0.00068,,,,,L
47374.32700000005,0,0.20896,0.00063,,,,,L
47374.46309999982,0,0.24149,0.00087,,,,,L
47394.28720000014,0,0.22347999999999998,0.0012900000000000001,,,,,L
47430.183699999936,0,0.26022,0.00102,,,,,L
47431.18839999987,0,0.2738,0.0011,,,,,L
47577.56059999997,0,0.24764,0.0012900000000000001,,,,,L
47710.376199999824,0,0.23207,0.0012900000000000001,,,,,L
47846.20679999981,0,0.20173,0.00103,,,,,L
48018.48950000014,0,0.19107,0.00164,,,,,L
48019.46309999982,0,0.21691,0.00163,,,,,L
48113.35319999978,0,0.15019,0.00102,,,,,L
48200.11340000015,0,0.17647,0.00103,,,,,L
48374.509300000034,0,0.15868000000000002,0.00163,,,,,L
48437.46029999992,0,0.1732,0.0015300000000000001,,,,,L
48834.32940000016,0,0.13343,0.00075,,,,,L
48847.32530000014,0,0.12849000000000002,0.00093,,,,,L
48905.20999999996,0,0.13685,0.0008399999999999999,,,,,L
49122.503800000064,0,0.09109,0.00138,,,,,L
49171.33900000015,0,0.10359,0.00111,,,,,L
49173.33939999994,0,0.09898,0.00126,,,,,L
49174.40260000015,0,0.11075,0.0011200000000000001,,,,,L
49200.32750000013,0,0.10012,0.0010400000000000001,,,,,L
49278.23779999977,0,0.105,0.00108,,,,,L
49280.248300000094,0,0.09792000000000001,0.00099,,,,,L
49587.26360000018,0,0.08056,0.00098,,,,,L
49589.26140000019,0,0.08245000000000001,0.00103,,,,,L
49602.246100000106,0,0.10185,0.0009599999999999999,,,,,L
49622.22470000014,0,0.09093000000000001,0.00126,,,,,L
49680.09459999995,0,0.06684,0.00076,,,,,L
49872.432200000156,0,0.05442,0.00049,,,,,L
49893.37559999991,0,0.06143,0.0005,,,,,L
49913.41770000011,0,0.045840000000000006,0.00051,,,,,L
49984.184599999804,0,0.056549999999999996,0.00051,,,,,L
50215.36720000021,0,0.00885,0.00078,,,,,L
50263.349200000055,0,0.02616,0.00054,,,,,L
50299.328600000124,0,0.03372,0.00061,,,,,L
50304.34679999994,0,0.01814,0.0005600000000000001,,,,,L
50601.384899999946,0,0.0,1e-05,,,,,K
50602.49110000022,0,0.00378,0.0012900000000000001,,,,,K
50603.521699999925,0,0.00369,0.0011899999999999999,,,,,K
50604.529800000135,0,0.00292,0.0013,,,,,K
50605.562100000214,0,-0.005730000000000001,0.0013,,,,,K
50606.46579999989,0,0.0007700000000000001,0.00138,,,,,K
50607.41569999978,0,0.00379,0.0013700000000000001,,,,,K
50613.400899999775,0,-0.00226,0.00079,,,,,L
50614.43129999982,0,-0.00037,0.00075,,,,,L
50640.39979999978,0,0.00626,0.00075,,,,,L
50656.32690000022,0,0.00741,0.00076,,,,,L
50793.062799999956,0,0.00176,0.00051,,,,,L
50979.42780000018,0,-0.00621,0.00116,,,,,L
51048.34739999985,0,0.00592,0.00117,,,,,L
51062.268999999855,0,-0.009130000000000001,0.00123,,,,,L
51068.345000000205,0,-0.023710000000000002,0.00145,,,,,K
51069.3909,0,-0.0292,0.00134,,,,,K
51070.41080000019,0,-0.019719999999999998,0.00131,,,,,K
51071.3470999999,0,-0.01315,0.0016,,,,,K
51072.33520000009,0,-0.01516,0.00134,,,,,K
51074.35490000015,0,-0.02296,0.00131,,,,,K
51075.28480000002,0,-0.0223,0.0011799999999999998,,,,,K
51299.5046000001,0,-0.0332,0.00078,,,,,L
51303.47779999999,0,-0.03321,0.0008,,,,,L
51305.47300000023,0,-0.03272,0.00088,,,,,L
51310.60020000022,0,-0.04188,0.00123,,,,,K
51311.5885999999,0,-0.040530000000000004,0.0013700000000000001,,,,,K
51312.601900000125,0,-0.033600000000000005,0.00139,,,,,K
51313.624799999874,0,-0.03907,0.0010400000000000001,,,,,K
51367.40590000013,0,-0.05311,0.00151,,,,,K
51368.37099999981,0,-0.049159999999999995,0.00172,,,,,K
51369.51519999979,0,-0.045700000000000005,0.00128,,,,,K
51371.52680000011,0,-0.041960000000000004,0.00164,,,,,K
51372.562299999874,0,-0.04615,0.00159,,,,,K
51373.3199,0,-0.0511,0.00148,,,,,K
51389.36549999984,0,-0.05127,0.0008,,,,,L
51392.33059999999,0,-0.0425,0.0008900000000000001,,,,,L
51410.376199999824,0,-0.05299,0.00152,,,,,K
51411.37170000002,0,-0.05586,0.00139,,,,,K
51416.203600000124,0,-0.02917,0.00088,,,,,L
51439.32209999999,0,-0.05396,0.00146,,,,,K
51467.139700000174,0,-0.028579999999999998,0.00157,,,,,L
51468.10719999997,0,-0.03472,0.00131,,,,,L
51678.5915000001,0,-0.03986,0.00185,,,,,K
51702.57480000006,0,-0.04632,0.0015,,,,,K
51751.32109999983,0,-0.05003,0.00088,,,,,L
51754.441800000146,0,-0.049210000000000004,0.00133,,,,,K
51793.325000000186,0,-0.053329999999999995,0.00163,,,,,K
51815.14759999979,0,-0.045259999999999995,0.00148,,,,,L
52030.53979999991,0,-0.07977,0.0015300000000000001,,,,,K
52041.49459999986,0,-0.07895999999999999,0.00122,,,,,L
52075.42989999987,0,-0.08123000000000001,0.00098,,,,,L
52098.54829999991,0,-0.08746,0.0016,,,,,K
52115.385799999814,0,-0.08766,0.00095,,,,,L
52117.36409999989,0,-0.08270999999999999,0.00082,,,,,L
52127.48309999984,0,-0.10094,0.00183,,,,,K
52133.30019999994,0,-0.09858,0.0016200000000000001,,,,,K
52140.274399999995,0,-0.0695,0.00087,,,,,L
52181.1483,0,-0.08426,0.00099,,,,,L
52186.168699999806,0,-0.07565999999999999,0.00115,,,,,L
52188.200900000054,0,-0.07717,0.0012900000000000001,,,,,L
52202.17659999989,0,-0.0917,0.00109,,,,,L
52390.63920000009,0,-0.12049,0.00191,,,,,K
52508.25889999978,0,-0.12032,0.0011799999999999998,,,,,L
52509.246199999936,0,-0.13404,0.00114,,,,,L
52515.282500000205,0,-0.11541,0.0011200000000000001,,,,,K
52575.19119999977,0,-0.11679,0.0016899999999999999,,,,,K
52848.40310000023,0,-0.13174,0.0017800000000000001,,,,,K
52855.498800000176,0,-0.13641999999999999,0.0016699999999999998,,,,,K
52894.229700000025,0,-0.13906,0.00145,,,,,L
53180.527199999895,0,-0.17106,0.00148,,,,,K
53203.32810000004,0,-0.16071000000000002,0.0019,,,,,L
53303.21130000008,0,-0.17217,0.0008399999999999999,,,,,K
53544.44480000017,0,-0.20304,0.00147,,,,,L
53550.50810000021,0,-0.19505,0.00055,,,,,K
53556.41270000022,0,-0.19072,0.0013700000000000001,,,,,L
53561.3936999999,0,-0.20113,0.00158,,,,,L
53589.3254999998,0,-0.1869,0.00281,,,,,L
53603.42050000001,0,-0.19575,0.00074,,,,,K
53925.53860000009,0,-0.21209999999999998,0.00067,,,,,K
53926.39570000023,0,-0.21761000000000003,0.00165,,,,,L
53954.259300000034,0,-0.22625,0.00197,,,,,L
53958.28280000016,0,-0.22132,0.0012900000000000001,,,,,L
53982.31380000012,0,-0.22035,0.00062,,,,,K
53989.37000000011,0,-0.22479,0.00304,,,,,L
54336.54339999985,0,-0.23541,0.0014399999999999999,,,,,K
54545.65180000011,0,-0.24005,0.00122,,,,,K
54671.45869999984,0,-0.2409,0.00121,,,,,K
54673.46560000023,0,-0.24816999999999997,0.00116,,,,,K
54675.38409999991,0,-0.24746,0.00093,,,,,K
54688.37419999996,0,-0.23927,0.00147,,,,,K
54689.45190000022,0,-0.2419,0.00136,,,,,K
54697.41969999997,0,-0.22274000000000002,0.00294,,,,,L
54717.43069999991,0,-0.24481,0.00125,,,,,K
54718.48639999982,0,-0.24641,0.00127,,,,,K
54719.49110000022,0,-0.23901,0.0011799999999999998,,,,,K
54720.45660000015,0,-0.22502,0.00122,,,,,K
54721.470100000035,0,-0.21899000000000002,0.0011899999999999999,,,,,K
54722.36959999986,0,-0.22622,0.0012,,,,,K
54724.426299999934,0,-0.23332,0.00127,,,,,K
54725.34159999993,0,-0.23283,0.00123,,,,,K
54726.45809999993,0,-0.23117,0.00116,,,,,K
54727.33650000021,0,-0.2365,0.00121,,,,,K
54777.31589999981,0,-0.2358,0.0014299999999999998,,,,,K
54808.20679999981,0,-0.2346,0.00135,,,,,K
54928.62869999977,0,-0.2344,0.0014299999999999998,,,,,K
54929.622099999804,0,-0.23722,0.00134,,,,,K
54934.629499999806,0,-0.23098,0.00122,,,,,K
55041.37009999994,0,-0.22955,0.00134,,,,,K
55076.242999999784,0,-0.23053,0.0012900000000000001,,,,,K
55079.23719999986,0,-0.23251,0.00134,,,,,K
55106.407600000035,0,-0.21921000000000002,0.00136,,,,,K
55169.256699999794,0,-0.2303,0.00088,,,,,K
55289.6483,0,-0.21821000000000002,0.00136,,,,,K
55312.63689999981,0,-0.22411,0.00132,,,,,K
55318.59970000014,0,-0.22652,0.00148,,,,,K
55404.40950000007,0,-0.21463,0.00135,,,,,K
55436.50959999999,0,-0.20643,0.00135,,,,,K
55542.1875,0,-0.19894,0.00115,,,,,K
55722.51220000023,0,-0.16358,0.00138,,,,,K
55785.46659999993,0,-0.16926,0.00158,,,,,K
55842.296300000045,0,-0.15836,0.00157,,,,,K
52143.5,1,,,786.0,6.0,157.9,0.5,
52253.5,1,,,794.0,5.0,157.3,0.6,
52472.5,1,,,788.0,6.0,156.6,0.9,
53988.5,1,,,750.0,80.0,155.0,5.0,
54636.5,1,,,742.0,35.0,151.8,2.9,
55696.5,1,,,519.0,6.0,147.1,0.5,
97 changes: 74 additions & 23 deletions orbitize/results.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import warnings
import h5py
import copy
import pdb
import itertools

import astropy.units as u
import astropy.constants as consts
Expand All @@ -17,7 +17,6 @@
import pandas as pd

import corner
import pdb

import orbitize.kepler as kepler
import orbitize.system
Expand Down Expand Up @@ -366,8 +365,8 @@ def plot_orbits(self, 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,
sep_pa_color='lightgrey', sep_pa_end_year=2025.0,
cbar_param='Epoch [year]', mod180=False, rv_time_series=False,plot_astrometry=True,
fig=None):
cbar_param='Epoch [year]', mod180=False, rv_time_series=False, plot_astrometry=True,
plot_astrometry_insts=False, fig=None):
"""
Plots one orbital period for a select number of fitted orbits
for a given object, with line segments colored according to time
Expand All @@ -394,7 +393,8 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
arcs with PAs that cross 360 deg during observations (default: False)
rv_time_series (Boolean): if fitting for secondary mass using MCMC for rv fitting and want to
display time series, set to True.
astrometry (Boolean): set to True by default. Plots the astrometric data.
plot_astrometry (Boolean): set to True by default. Plots the astrometric data.
plot_astrometry_insts (Boolean): set to False by default. Plots the astrometric data by instruments.
fig (matplotlib.pyplot.Figure): optionally include a predefined Figure object to plot the orbit on.
Most users will not need this keyword.
Expand All @@ -406,7 +406,6 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
Additions by Malena Rice, 2019
"""

if Time(start_mjd, format='mjd').decimalyear >= sep_pa_end_year:
raise ValueError('start_mjd keyword date must be less than sep_pa_end_year keyword date.')

Expand All @@ -419,6 +418,7 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
with warnings.catch_warnings():
warnings.simplefilter('ignore', ErfaWarning)

data = self.data
dict_of_indices = {
'sma': 0,
'ecc': 1,
Expand All @@ -442,7 +442,6 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
raise Exception(
'Invalid input; acceptable inputs include epochs, sma1, ecc1, inc1, aop1, pan1, tau1, sma2, ecc2, ...')


start_index = (object_to_plot - 1) * 6

sma = self.post[:, start_index + dict_of_indices['sma']]
Expand Down Expand Up @@ -507,6 +506,12 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
vmax=np.max(Time(epochs, format='mjd').decimalyear)
)

# Before starting to plot rv data, make sure rv data exists:
rv_indices = np.where(data['quant_type'] == 'rv')
if rv_time_series and len(rv_indices) == 0:
warnings.warn("Unable to plot radial velocity data.")
rv_time_series = False

# Create figure for orbit plots
if fig is None:
fig = plt.figure(figsize=(14, 6))
Expand All @@ -523,7 +528,6 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
else:
ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6)

data=self.data
astr_inds=np.where((~np.isnan(data['quant1'])) & (~np.isnan(data['quant2'])))
astr_epochs=data['epoch'][astr_inds]

Expand Down Expand Up @@ -557,6 +561,22 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
pa_data = np.append(pa_data, pa_from_dec_data)
pa_err = np.append(pa_err, pa_err_from_dec_data)

# For plotting different astrometry instruments
if plot_astrometry_insts:
astr_colors = ('#FF7F11', '#11FFE3', '#14FF11', '#7A11FF', '#FF1919')
astr_symbols = ('*', 'o', 'p', 's')

ax_colors = itertools.cycle(astr_colors)
ax_symbols = itertools.cycle(astr_symbols)

astr_data = data[astr_inds]
astr_insts = np.unique(data[astr_inds]['instrument'])

# Indices corresponding to each instrument in datafile
astr_inst_inds = {}
for i in range(len(astr_insts)):
astr_inst_inds[astr_insts[i]]=np.where(astr_data['instrument']==astr_insts[i].encode())[0]

# 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)
Expand All @@ -572,7 +592,16 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,

if plot_astrometry:
ra_data,dec_data=orbitize.system.seppa2radec(sep_data,pa_data)
ax.scatter(ra_data,dec_data,marker='*',c='#FF7F11',zorder=10,s=60)

# Plot astrometry along with instruments
if plot_astrometry_insts:
for i in range(len(astr_insts)):
ra = ra_data[astr_inst_inds[astr_insts[i]]]
dec = dec_data[astr_inst_inds[astr_insts[i]]]
ax.scatter(ra, dec, marker=next(ax_symbols), c=next(ax_colors), zorder=10, s=60, label=astr_insts[i])
else:
ax.scatter(ra_data, dec_data, marker='*', c='#FF7F11', zorder=10, s=60)

# modify the axes
if square_plot:
adjustable_param = 'datalim'
Expand Down Expand Up @@ -605,6 +634,13 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
ax1.set_ylabel('$\\rho$ [mas]')
ax2.set_xlabel('Epoch')

if plot_astrometry_insts:
ax1_colors = itertools.cycle(astr_colors)
ax1_symbols = itertools.cycle(astr_symbols)

ax2_colors = itertools.cycle(astr_colors)
ax2_symbols = itertools.cycle(astr_symbols)

epochs_seppa = np.zeros((num_orbits_to_plot, num_epochs_to_plot))

for i in np.arange(num_orbits_to_plot):
Expand Down Expand Up @@ -642,23 +678,35 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,

plt.sca(ax1)
plt.plot(yr_epochs, seps, color=sep_pa_color)
# plot separations from data points
plt.scatter(Time(astr_epochs,format='mjd').decimalyear,sep_data,s=10,marker='*',c='purple',zorder=10)

plt.sca(ax2)
plt.plot(yr_epochs, pas, color=sep_pa_color)

# Plot sep/pa instruments
if plot_astrometry_insts:
for i in range(len(astr_insts)):
sep = sep_data[astr_inst_inds[astr_insts[i]]]
pa = pa_data[astr_inst_inds[astr_insts[i]]]
epochs = astr_epochs[astr_inst_inds[astr_insts[i]]]
plt.sca(ax1)
plt.scatter(Time(epochs,format='mjd').decimalyear,sep,s=10,marker=next(ax1_symbols),c=next(ax1_colors),zorder=10,label=astr_insts[i])
plt.sca(ax2)
plt.scatter(Time(epochs,format='mjd').decimalyear,pa,s=10,marker=next(ax2_symbols),c=next(ax2_colors),zorder=10)
plt.sca(ax1)
plt.legend(title='Instruments', bbox_to_anchor=(1.3, 1), loc='upper right')
else:
plt.sca(ax1)
plt.scatter(Time(astr_epochs,format='mjd').decimalyear,sep_data,s=10,marker='*',c='purple',zorder=10)
plt.sca(ax2)
plt.scatter(Time(astr_epochs,format='mjd').decimalyear,pa_data,s=10,marker='*',c='purple',zorder=10)

if rv_time_series:

# switch current axis to rv panel
plt.sca(ax3)

# get list of instruments
insts=np.unique(data['instrument'])
insts=[i if isinstance(i,str) else i.decode() for i in insts]
insts=[i for i in insts if 'def' not in i]

# get list of rv instruments
insts = np.unique(data['instrument'][rv_indices])

# get gamma/sigma labels and corresponding positions in the posterior
gams=['gamma_'+inst for inst in insts]

Expand All @@ -676,12 +724,15 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
inds[insts[i]]=np.where(data['instrument']==insts[i].encode())[0]

# choose the orbit with the best log probability
best_like=np.where(self.lnlike==np.amin(self.lnlike))[0][0]
best_like=np.where(self.lnlike==np.amax(self.lnlike))[0][0]
med_ga=[self.post[best_like,i] for i in gam_idx]

# colour/shape scheme scheme for rv data points
clrs=['0496FF','372554','FF1053','3A7CA5','143109']
symbols=['o','^','v','s']
clrs=('#0496FF','#372554','#FF1053','#3A7CA5','#143109')
symbols=('o','^','v','s')

ax3_colors = itertools.cycle(clrs)
ax3_symbols = itertools.cycle(symbols)

# get rvs and plot them
for i,name in enumerate(inds.keys()):
Expand All @@ -691,15 +742,15 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
epochs=inst_data['epoch']
epochs=Time(epochs, format='mjd').decimalyear
rvs-=med_ga[i]
plt.scatter(epochs,rvs,marker=symbols[i],s=5,label=name,c=f'#{clrs[i]}',zorder=5)
plt.scatter(epochs,rvs,s=5,marker=next(ax3_symbols),c=next(ax3_colors),label=name,zorder=5)

inds[insts[i]]=np.where(data['instrument']==insts[i])[0]
plt.legend()


# calculate the predicted rv trend using the best orbit
raa, decc, vz = kepler.calc_orbit(
epochs_seppa[i, :], sma[best_like], ecc[best_like], inc[best_like], aop[best_like], pan[best_like],
epochs_seppa[0, :], sma[best_like], ecc[best_like], inc[best_like], aop[best_like], pan[best_like],
tau[best_like], plx[best_like], mtot[best_like], tau_ref_epoch=self.tau_ref_epoch,
mass_for_Kamp=m0[best_like]
)
Expand All @@ -708,7 +759,7 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,

# plot rv trend

plt.plot(Time(epochs_seppa[i, :],format='mjd').decimalyear, vz, color=sep_pa_color)
plt.plot(Time(epochs_seppa[0, :],format='mjd').decimalyear, vz, color=sep_pa_color)


# add colorbar
Expand Down

0 comments on commit 44f86a5

Please sign in to comment.