import numpy as np
import corner
import warnings
import itertools
import astropy.units as u
import astropy.constants as consts
from astropy.time import Time
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import matplotlib.colors as colors
from erfa import ErfaWarning
import orbitize
import orbitize.kepler as kepler
# TODO: deprecatation warning for plots in results
# define modified color map for default use in orbit plots
cmap = mpl.cm.Purples_r
cmap = colors.LinearSegmentedColormap.from_list(
'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=0.0, b=0.7),
cmap(np.linspace(0.0, 0.7, 1000))
)
[docs]def plot_corner(results, param_list=None, **corner_kwargs):
"""
Make a corner plot of posterior on orbit fit from any sampler
Args:
param_list (list of strings): each entry is a name of a parameter to include.
Valid strings::
sma1: semimajor axis
ecc1: eccentricity
inc1: inclination
aop1: argument of periastron
pan1: position angle of nodes
tau1: epoch of periastron passage, expressed as fraction of orbital period
per1: period
K1: stellar radial velocity semi-amplitude
[repeat for 2, 3, 4, etc if multiple objects]
plx: parallax
pm_ra: RA proper motion
pm_dec: Dec proper motion
alpha0: primary offset from reported Hipparcos RA @ alphadec0_epoch (generally 1991.25)
delta0: primary offset from reported Hipparcos Dec @ alphadec0_epoch (generally 1991.25)
gamma: rv offset
sigma: rv jitter
mi: mass of individual body i, for i = 0, 1, 2, ... (only if fit_secondary_mass == True)
mtot: total mass (only if fit_secondary_mass == False)
**corner_kwargs: any remaining keyword args are sent to ``corner.corner``.
See `here <https://corner.readthedocs.io/>`_.
Note: default axis labels used unless overwritten by user input.
Return:
``matplotlib.pyplot.Figure``: corner plot
.. Note:: **Example**: Use ``param_list = ['sma1,ecc1,inc1,sma2,ecc2,inc2']`` to only
plot posteriors for semimajor axis, eccentricity and inclination
of the first two companions
Written: Henry Ngo, 2018
"""
# Define array of default axis labels (overwritten if user specifies list)
default_labels = {
'sma': '$a_{0}$ [au]',
'ecc': '$ecc_{0}$',
'inc': '$inc_{0}$ [$^\\circ$]',
'aop': '$\\omega_{0}$ [$^\\circ$]',
'pan': '$\\Omega_{0}$ [$^\\circ$]',
'tau': '$\\tau_{0}$',
'plx': '$\\pi$ [mas]',
'gam': '$\\gamma$ [km/s]',
'sig': '$\\sigma$ [km/s]',
'mtot': '$M_T$ [M$_{{\\odot}}$]',
'm0': '$M_0$ [M$_{{\\odot}}$]',
'm': '$M_{0}$ [M$_{{\\rm Jup}}$]',
'pm_ra': '$\\mu_{{\\alpha}}$ [mas/yr]',
'pm_dec': '$\\mu_{{\\delta}}$ [mas/yr]',
'alpha0': '$\\alpha^{{*}}_{{0}}$ [mas]',
'delta0': '$\\delta_0$ [mas]',
'm': '$M_{0}$ [M$_{{\\rm Jup}}$]',
'per' : '$P_{0}$ [yr]',
'K' : '$K_{0}$ [km/s]',
'x' : '$X_{0}$ [AU]',
'y' : '$Y_{0}$ [AU]',
'z' : '$Z_{0}$ [AU]',
'xdot' : '$xdot_{0}$ [km/s]',
'ydot' : '$ydot_{0}$ [km/s]',
'zdot' : '$zdot_{0}$ [km/s]'
}
if param_list is None:
param_list = results.labels
param_indices = []
angle_indices = []
secondary_mass_indices = []
for i, param in enumerate(param_list):
index_num = results.param_idx[param]
# only plot non-fixed parameters
if np.std(results.post[:, index_num]) > 0:
param_indices.append(index_num)
label_key = param
if label_key.startswith('aop') or label_key.startswith('pan') or label_key.startswith('inc'):
angle_indices.append(i)
if label_key.startswith('m') and label_key != 'm0' and label_key != 'mtot':
secondary_mass_indices.append(i)
samples = np.copy(results.post[:, param_indices]) # keep only chains for selected parameters
samples[:, angle_indices] = np.degrees(
samples[:, angle_indices]) # convert angles from rad to deg
samples[:, secondary_mass_indices] *= u.solMass.to(u.jupiterMass) # convert to Jupiter masses for companions
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]
if label_key.startswith("m") and label_key != 'm0' and label_key != 'mtot':
body_num = label_key[1]
label_key = "m"
elif label_key == 'm0' or label_key == 'mtot' or label_key.startswith('plx'):
body_num = ""
# maintain original label key
elif label_key in ['pm_ra', 'pm_dec', 'alpha0', 'delta0']:
body_num = ""
elif label_key.startswith("gamma") or label_key.startswith("sigma"):
body_num = ""
label_key = label_key[0:3]
else:
body_num = label_key[-1]
label_key = label_key[0:-1]
reduced_labels_list.append(default_labels[label_key].format(body_num))
corner_kwargs['labels'] = reduced_labels_list
figure = corner.corner(samples, **corner_kwargs)
return figure
[docs]def plot_orbits(results, 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,
plot_astrometry_insts=False, plot_errorbars=True, fig=None):
"""
Plots one orbital period for a select number of fitted orbits
for a given object, with line segments colored according to time
Args:
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)
square_plot (Boolean): Aspect ratio is always equal, but if
square_plot is True (default), then the axes will be square,
otherwise, white space padding is used
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).
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)
rv_time_series (Boolean): if fitting for secondary mass using MCMC for rv fitting and want to
display time series, set to True.
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.
plot_errorbars (Boolean): set to True by default. Plots error bars of measurements
fig (matplotlib.pyplot.Figure): optionally include a predefined Figure object to plot the orbit on.
Most users will not need this keyword.
Return:
``matplotlib.pyplot.Figure``: the orbit plot if input is valid, ``None`` otherwise
(written): Henry Ngo, Sarah Blunt, 2018
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.')
if object_to_plot > results.num_secondary_bodies:
raise ValueError("Only {0} secondary bodies being fit. Requested to plot body {1} which is out of range".format(results.num_secondary_bodies, object_to_plot))
if object_to_plot == 0:
raise ValueError("Plotting the primary's orbit is currently unsupported. Stay tuned.")
if rv_time_series and 'm0' not in results.labels:
rv_time_series = False
warnings.warn("It seems that the stellar and companion mass "
"have not been fitted separately. Setting "
"rv_time_series=True is therefore not possible "
"so the argument is set to False instead.")
with warnings.catch_warnings():
warnings.simplefilter('ignore', ErfaWarning)
data = results.data[results.data['object'] == object_to_plot]
possible_cbar_params = [
'sma',
'ecc',
'inc',
'aop'
'pan',
'tau',
'plx'
]
if cbar_param == 'Epoch [year]':
pass
elif cbar_param[0:3] in possible_cbar_params:
index = results.param_idx[cbar_param]
else:
raise Exception(
"Invalid input; acceptable inputs include 'Epoch [year]', 'plx', 'sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'sma2', 'ecc2', ...)"
)
# Select random indices for plotted orbit
num_orbits = len(results.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)
# Get posteriors from random indices
standard_post = []
if results.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(results.post[orb_ind])
standard_post.append(results.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(results.post[orb_ind])
standard_post = np.array(standard_post)
sma = standard_post[:, results.standard_param_idx['sma{}'.format(object_to_plot)]]
ecc = standard_post[:, results.standard_param_idx['ecc{}'.format(object_to_plot)]]
inc = standard_post[:, results.standard_param_idx['inc{}'.format(object_to_plot)]]
aop = standard_post[:, results.standard_param_idx['aop{}'.format(object_to_plot)]]
pan = standard_post[:, results.standard_param_idx['pan{}'.format(object_to_plot)]]
tau = standard_post[:, results.standard_param_idx['tau{}'.format(object_to_plot)]]
plx = standard_post[:, results.standard_param_idx['plx']]
# Then, get the other parameters
if 'mtot' in results.labels:
mtot = standard_post[:, results.standard_param_idx['mtot']]
elif 'm0' in results.labels:
m0 = standard_post[:, results.standard_param_idx['m0']]
m1 = standard_post[:, results.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))
# Loop through each orbit to plot and calcualte ra/dec offsets for all points in orbit
# Need this loops since epochs[] vary for each orbit, unless we want to just plot the same time period for all orbits
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(start_mjd, float(
start_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=results.tau_ref_epoch
)
raoff[i, :] = raoff0
deoff[i, :] = deoff0
# Create a linearly increasing colormap for our range of epochs
if cbar_param != 'Epoch [year]':
cbar_param_arr = results.post[:, index]
norm = mpl.colors.Normalize(vmin=np.min(cbar_param_arr),
vmax=np.max(cbar_param_arr))
norm_yr = mpl.colors.Normalize(vmin=np.min(
cbar_param_arr), vmax=np.max(cbar_param_arr))
elif cbar_param == 'Epoch [year]':
min_cbar_date = np.min(epochs)
max_cbar_date = np.max(epochs[-1, :])
# if we're plotting orbital periods greater than 1,000 yrs, limit the colorbar dynamic range
if max_cbar_date - min_cbar_date > 1000 * 365.25:
max_cbar_date = min_cbar_date + 1000 * 365.25
norm = mpl.colors.Normalize(vmin=min_cbar_date, vmax=max_cbar_date)
norm_yr = mpl.colors.Normalize(
vmin=Time(min_cbar_date, format='mjd').decimalyear,
vmax=Time(max_cbar_date, 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))
if rv_time_series:
fig = plt.figure(figsize=(14, 9))
ax = plt.subplot2grid((3, 14), (0, 0), rowspan=2, colspan=6)
else:
fig = plt.figure(figsize=(14, 6))
ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6)
else:
plt.set_current_figure(fig)
if rv_time_series:
ax = plt.subplot2grid((3, 14), (0, 0), rowspan=2, colspan=6)
else:
ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6)
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
# For plotting different astrometry instruments
if plot_astrometry_insts:
astr_colors = ('purple','#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)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(
segments, cmap=cmap, norm=norm, linewidth=1.0
)
if cbar_param != 'Epoch [year]':
lc.set_array(np.ones(len(epochs[0]))*cbar_param_arr[i])
elif cbar_param == 'Epoch [year]':
lc.set_array(epochs[i, :])
ax.add_collection(lc)
# if plot_astrometry:
# # 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]]]
# if plot_errorbars:
# xerr = ra_err[astr_inst_inds[astr_insts[i]]]
# yerr = dec_err[astr_inst_inds[astr_insts[i]]]
# else:
# xerr = None
# yerr = None
# ax.errorbar(ra, dec, xerr=xerr, yerr=yerr, marker=next(ax_symbols), c=next(ax_colors), zorder=10, label=astr_insts[i], linestyle='', ms=5, capsize=2)
# else:
# if plot_errorbars:
# xerr = ra_err
# yerr = dec_err
# else:
# xerr = None
# yerr = None
# ax.errorbar(ra_data, dec_data, xerr=xerr, yerr=yerr, marker='o', c='#FF7F11', zorder=10, linestyle='', capsize=2, ms=5)
# modify the axes
if square_plot:
adjustable_param = 'datalim'
else:
adjustable_param = 'box'
ax.set_aspect('equal', adjustable=adjustable_param)
ax.set_xlabel('$\\Delta$RA [mas]')
ax.set_ylabel('$\\Delta$Dec [mas]')
ax.locator_params(axis='x', nbins=6)
ax.locator_params(axis='y', nbins=6)
ax.invert_xaxis() # To go to a left-handed coordinate system
# plot sep/PA and/or rv zoom-in panels
if rv_time_series:
ax1 = plt.subplot2grid((3, 14), (0, 8), colspan=6)
ax2 = plt.subplot2grid((3, 14), (1, 8), colspan=6)
ax3 = plt.subplot2grid((3, 14), (2, 0), colspan=14, rowspan=1)
ax2.set_ylabel('PA [$^{{\\circ}}$]')
ax1.set_ylabel('$\\rho$ [mas]')
ax3.set_ylabel('RV [km/s]')
ax3.set_xlabel('Epoch')
ax2.set_xlabel('Epoch')
plt.subplots_adjust(hspace=0.3)
else:
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')
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):
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
if rv_time_series:
raoff0, deoff0, _ = kepler.calc_orbit(
epochs_seppa[i, :], sma[i], ecc[i], inc[i], aop[i], pan[i],
tau[i], plx[i], mtot[i], tau_ref_epoch=results.tau_ref_epoch,
mass_for_Kamp=m0[i]
)
raoff[i, :] = raoff0
deoff[i, :] = deoff0
else:
raoff0, deoff0, _ = kepler.calc_orbit(
epochs_seppa[i, :], sma[i], ecc[i], inc[i], aop[i], pan[i],
tau[i], plx[i], mtot[i], tau_ref_epoch=results.tau_ref_epoch
)
raoff[i, :] = raoff0
deoff[i, :] = deoff0
yr_epochs = Time(epochs_seppa[i, :], format='mjd').decimalyear
seps, pas = orbitize.system.radec2seppa(raoff[i, :], deoff[i, :], mod180=mod180)
plt.sca(ax1)
plt.plot(yr_epochs, seps, color=sep_pa_color, zorder=1)
plt.sca(ax2)
plt.plot(yr_epochs, pas, color=sep_pa_color, zorder=1)
# 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]]]
if plot_errorbars:
serr = sep_err[astr_inst_inds[astr_insts[i]]]
perr = pa_err[astr_inst_inds[astr_insts[i]]]
else:
yerr = None
perr = None
plt.sca(ax1)
plt.errorbar(Time(epochs,format='mjd').decimalyear,sep,yerr=serr,ms=5, linestyle='',marker=next(ax1_symbols),c=next(ax1_colors),zorder=10,label=astr_insts[i], capsize=2)
plt.sca(ax2)
plt.errorbar(Time(epochs,format='mjd').decimalyear,pa,yerr=perr,ms=5, linestyle='',marker=next(ax2_symbols),c=next(ax2_colors),zorder=10, capsize=2)
plt.sca(ax1)
plt.legend(title='Instruments', bbox_to_anchor=(1.3, 1), loc='upper right')
else:
if plot_errorbars:
serr = sep_err
perr = pa_err
else:
yerr = None
perr = None
plt.sca(ax1)
plt.errorbar(Time(astr_epochs,format='mjd').decimalyear,sep_data,yerr=serr,ms=5, linestyle='',marker='o',c='purple',zorder=2, capsize=2)
plt.sca(ax2)
plt.errorbar(Time(astr_epochs,format='mjd').decimalyear,pa_data,yerr=perr,ms=5, linestyle='',marker='o',c='purple',zorder=2, capsize=2)
if rv_time_series:
rv_data = results.data[results.data['object'] == 0]
rv_data = rv_data[rv_data['quant_type'] == 'rv']
# switch current axis to rv panel
plt.sca(ax3)
# get list of rv instruments
insts = np.unique(rv_data['instrument'])
if len(insts) == 0:
insts = ['defrv']
# get gamma/sigma labels and corresponding positions in the posterior
gams=['gamma_'+inst for inst in insts]
sigs = ['sigma_'+inst for inst in insts]
if isinstance(results.labels,list):
labels=np.array(results.labels)
else:
labels=results.labels
# get the indices corresponding to each gamma within results.labels
gam_idx=[np.where(labels==inst_gamma)[0] for inst_gamma in gams]
# indices corresponding to each instrument in the datafile
inds={}
for i in range(len(insts)):
inds[insts[i]]=np.where(rv_data['instrument']==insts[i].encode())[0]
# choose the orbit with the best log probability
best_like=np.where(results.lnlike==np.amax(results.lnlike))[0][0]
# Get the posteriors for this index and convert to standard basis
best_post = results.basis.to_standard_basis(results.post[best_like].copy())
# Get the masses for the best posteriors:
best_m0 = best_post[results.standard_param_idx['m0']]
best_m1 = best_post[results.standard_param_idx['m{}'.format(object_to_plot)]]
best_mtot = best_m0 + best_m1
# colour/shape scheme scheme for rv data points
clrs=('purple', '#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()):
inst_data=rv_data[inds[name]]
rvs=inst_data['quant1']
epochs=inst_data['epoch']
epochs=Time(epochs, format='mjd').decimalyear
rvs -= best_post[results.param_idx[gams[i]]]
if plot_errorbars:
yerr = inst_data['quant1_err']
yerr = np.sqrt(yerr**2 + best_post[results.param_idx[sigs[i]]]**2)
plt.errorbar(epochs,rvs,yerr=yerr,ms=5, linestyle='',marker=next(ax3_symbols),c=next(ax3_colors),label=name,zorder=5,capsize=2)
if len(inds.keys()) == 1 and 'defrv' in inds.keys():
pass
else:
plt.legend()
# calculate the predicted rv trend using the best orbit
_, _, vz = kepler.calc_orbit(
epochs_seppa[0, :],
best_post[results.standard_param_idx['sma{}'.format(object_to_plot)]],
best_post[results.standard_param_idx['ecc{}'.format(object_to_plot)]],
best_post[results.standard_param_idx['inc{}'.format(object_to_plot)]],
best_post[results.standard_param_idx['aop{}'.format(object_to_plot)]],
best_post[results.standard_param_idx['pan{}'.format(object_to_plot)]],
best_post[results.standard_param_idx['tau{}'.format(object_to_plot)]],
best_post[results.standard_param_idx['plx']], best_mtot,
tau_ref_epoch=results.tau_ref_epoch, mass_for_Kamp=best_m0
)
vz=vz*-(best_m1)/np.median(best_m0)
# plot rv trend
plt.plot(Time(epochs_seppa[0, :],format='mjd').decimalyear, vz, color=sep_pa_color, zorder=1)
# add colorbar
if show_colorbar:
if rv_time_series:
# Create an axes for colorbar. The position of the axes is calculated based on the position of ax.
# You can change x1.0.05 to adjust the distance between the main image and the colorbar.
# You can change 0.02 to adjust the width of the colorbar.
cbar_ax = fig.add_axes(
[ax.get_position().x1+0.005, ax.get_position().y0, 0.02, ax.get_position().height])
cbar = mpl.colorbar.ColorbarBase(
cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param)
else:
# xpos, ypos, width, height, in fraction of figure size
cbar_ax = fig.add_axes([0.47, 0.15, 0.015, 0.7])
cbar = mpl.colorbar.ColorbarBase(
cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param)
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