def test_orbit_e99():
"""
Test a highly eccentric orbit (ecc=0.99). Again validate against James Graham's orbit code
"""
# sma, ecc, inc, argp, lan, tau, plx, mtot
orbital_params = np.array([10, 0.99, 3, 0.5, 1.5, 0.3, 50, 1.5])
epochs = np.array([1000, 1101.4])
raoffs, deoffs, vzs = kepler.calc_orbit(epochs,
orbital_params[0],
orbital_params[1],
orbital_params[2],
orbital_params[3],
orbital_params[4],
orbital_params[5],
orbital_params[6],
orbital_params[7],
tau_ref_epoch=0)
true_raoff = [-589.45575, -571.48432]
true_deoff = [-447.32217, -437.68456]
true_vz = [.39208876, .42041953]
for meas, truth in zip(raoffs, true_raoff):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(deoffs, true_deoff):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(vzs, true_vz):
assert truth == pytest.approx(meas, abs=1e-8)
def test_orbit_with_mass():
"""
Test a orbit where we specify the mass of the body too. This will change the radial velocity, which normally assumes the body is a test particle
We will test two equal mass bodies, which will reduce the RV signal by 2, compared to the RV signal of a massless particle in a system with the
same total mass.
"""
# sma, ecc, inc, argp, lan, tau, plx, mtot
orbital_params = np.array([10, 0.99, 3, 0.5, 1.5, 0.3, 50, 1.5])
epochs = np.array([1000, 1101.4])
raoffs, deoffs, vzs = kepler.calc_orbit(epochs,
orbital_params[0],
orbital_params[1],
orbital_params[2],
orbital_params[3],
orbital_params[4],
orbital_params[5],
orbital_params[6],
orbital_params[7],
mass_for_Kamp=orbital_params[7] /
2)
true_raoff = [-589.45575, -571.48432]
true_deoff = [-447.32217, -437.68456]
true_vz = [.39208876 / 2, .42041953 / 2]
for meas, truth in zip(raoffs, true_raoff):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(deoffs, true_deoff):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(vzs, true_vz):
assert truth == pytest.approx(meas, abs=1e-8)
def test_orbit_scalar():
"""
Test orbitize.kepler.calc_orbit() with scalar values
"""
sma = 10
ecc = 0.3
inc = 3
argp = 0.5
lan = 1.5
tau = 0.3
plx = 50
mtot = 1.5
epochs = 1000
raoffs, deoffs, vzs = kepler.calc_orbit(epochs,
sma,
ecc,
inc,
argp,
lan,
tau,
plx,
mtot,
tau_ref_epoch=0)
true_raoff = 152.86786
true_deoff = -462.91038
true_vz = .86448656
assert true_raoff == pytest.approx(raoffs, abs=threshold)
assert true_deoff == pytest.approx(deoffs, abs=threshold)
assert true_vz == pytest.approx(vzs, abs=1e-8)
def test_orbit_e03():
"""
Test orbitize.kepler.calc_orbit() by comparing this code to the output of James Graham's code which has been used in
many published papers. Note that orbitize currently uses Rob De Rosa's eccentricity solver.
Pretty standard orbit with ecc = 0.3
"""
# sma, ecc, inc, argp, lan, tau, plx, mtot
orbital_params = np.array([10, 0.3, 3, 0.5, 1.5, 0.3, 50, 1.5])
epochs = np.array([1000, 1101.4])
raoffs, deoffs, vzs = kepler.calc_orbit(
epochs, orbital_params[0], orbital_params[1], orbital_params[2],
orbital_params[3], orbital_params[4], orbital_params[5],
orbital_params[6], orbital_params[7])
true_raoff = [152.86786, 180.39408] #mas
true_deoff = [-462.91038, -442.0127]
true_vz = [.86448656, .97591289]
for meas, truth in zip(raoffs, true_raoff):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(deoffs, true_deoff):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(vzs, true_vz):
assert truth == pytest.approx(meas, abs=1e-8)
def test_orbit_e03_array():
"""
Test orbitize.kepler.calc_orbit() with a standard orbit with ecc = 0.3 and an array of keplerian input
"""
# sma, ecc, inc, argp, lan, tau, plx, mtot
sma = np.array([10, 10, 10])
ecc = np.array([0.3, 0.3, 0.3])
inc = np.array([3, 3, 3])
argp = np.array([0.5, 0.5, 0.5])
lan = np.array([1.5, 1.5, 1.5])
tau = np.array([0.3, 0.3, 0.3])
plx = np.array([50, 50, 50])
mtot = np.array([1.5, 1.5, 1.5])
epochs = np.array([1000, 1101.4])
raoffs, deoffs, vzs = kepler.calc_orbit(epochs, sma, ecc, inc, argp, lan,
tau, plx, mtot)
true_raoff = np.array([[152.86786, 152.86786, 152.86786],
[180.39408, 180.39408, 180.39408]])
true_deoff = np.array([[-462.91038, -462.91038, -462.91038],
[-442.0127, -442.0127, -442.0127]])
true_vz = np.array([[.86448656, .86448656, .86448656],
[.97591289, .97591289, .97591289]])
for ii in range(0, 3):
for meas, truth in zip(raoffs[:, ii], true_raoff[:, ii]):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(deoffs[:, ii], true_deoff[:, ii]):
assert truth == pytest.approx(meas, abs=threshold)
for meas, truth in zip(vzs[:, ii], true_vz[:, ii]):
assert truth == pytest.approx(meas, abs=1e-8)
def compute_model(self, params_arr):
"""
Compute model predictions for an array of fitting parameters.
Args:
params_arr (np.array of float): RxM array
of fitting parameters, where R is the number of
parameters being fit, and M is the number of orbits
we need model predictions for. Must be in the same order
documented in ``System()`` above. If M=1, this can be a 1d array.
Returns:
np.array of float: Nobsx2xM array model predictions. If M=1, this is
a 2d array, otherwise it is a 3d array.
"""
if len(params_arr.shape) == 1:
model = np.zeros((len(self.data_table), 2))
else:
model = np.zeros((len(self.data_table), 2, params_arr.shape[1]))
for body_num in np.arange(self.num_secondary_bodies)+1:
epochs = self.data_table['epoch'][self.body_indices[body_num]]
sma = params_arr[body_num-1]
ecc = params_arr[body_num]
inc = params_arr[body_num+1]
argp = params_arr[body_num+2]
lan = params_arr[body_num+3]
tau = params_arr[body_num+4]
plx = params_arr[6*self.num_secondary_bodies]
if self.fit_secondary_mass:
# mass of secondary bodies are in order from -1-num_bodies until -2 in order.
mass = params_arr[-1-self.num_secondary_bodies+(body_num-1)]
else:
mass = None
mtot = params_arr[-1]
raoff, decoff, vz = kepler.calc_orbit(
epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=mass, tau_ref_epoch=self.tau_ref_epoch
)
if len(raoff[self.radec[body_num]]) > 0: # (prevent empty array dimension errors)
model[self.radec[body_num], 0] = raoff[self.radec[body_num]]
model[self.radec[body_num], 1] = decoff[self.radec[body_num]]
if len(raoff[self.seppa[body_num]]) > 0:
sep, pa = radec2seppa(
raoff[self.seppa[body_num]],
decoff[self.seppa[body_num]]
)
model[self.seppa[body_num], 0] = sep
model[self.seppa[body_num], 1] = pa
return model
def gen_data():
'''
Simulates radial velocity and astrometric data for a test system.
Returns:
(rvs,rv_epochs): Tuple of generated radial velocity measurements (rvs) and their corresponding
measurement epochs (rv_epochs)
(ras,decs,astro_epochs): Tuple containing simulated astrometric measurements (ras, decs)
and the corresponding measurement epochs (astro_epochs)
'''
#set parameters for the synthetic data
sma=1
inc=np.pi/2
ecc=0.2
aop=np.pi/4
pan=np.pi/4
tau=0.4
plx=50
mass_for_kamp=0.1
mtot=1.1
#epochs and errors for rv
rv_epochs=np.linspace(51544,52426,200)
#epochs and errors for astrometry
astro_epochs=np.linspace(51500,52500,10)
astro_err=0
#generate rv trend
rvset=kepler.calc_orbit(rv_epochs,sma,ecc,inc,aop,pan,tau,plx,mtot,mass_for_Kamp=mass_for_kamp)
rvs=rvset[2]
#generate predictions for astrometric epochs
astro_set=kepler.calc_orbit(astro_epochs,sma,ecc,inc,aop,pan,tau,plx,mtot,mass_for_Kamp=mass_for_kamp)
ras,decs=astro_set[0],astro_set[1]
#return model generations
return (rvs,rv_epochs),(ras,decs,astro_epochs)
def test_2planet_nomass():
"""
Compare multiplanet to rebound for planets with mass.
"""
# generate a planet orbit
mjup = u.Mjup.to(u.Msun)
mass_b = 0 * mjup
mass_c = 0 * mjup
params = np.array([
10, 0.1,
np.radians(45),
np.radians(45),
np.radians(45), 0.5, 3, 0.1,
np.radians(45),
np.radians(190),
np.radians(45), 0.2, 1, mass_b, mass_c, 1.5 - mass_b - mass_c
])
tau_ref_epoch = 0
epochs = np.linspace(0, 365.25 * 4,
100) + tau_ref_epoch # nearly the full period, MJD
# doesn't matter that this is right, just needs to be the same size. below doesn't include effect of c
# just want to generate some measurements of plaent b to test compute model
b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit(
epochs,
params[0],
params[1],
params[2],
params[3],
params[4],
params[5],
params[-2],
params[-1],
tau_ref_epoch=tau_ref_epoch)
# generate some fake measurements of planet b, just to feed into system.py to test bookkeeping
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int), b_ra_model,
np.zeros(b_ra_model.shape), b_dec_model,
np.zeros(b_dec_model.shape)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
filename = os.path.join(orbitize.DATADIR, "rebound_2planet.csv")
t.write(filename)
# create the orbitize system and generate model predictions using the ground truth
astrom_dat = read_input.read_file(filename)
sys = system.System(2,
astrom_dat,
params[-1],
params[-2],
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
# generate measurement
radec_orbitize, _ = sys.compute_model(params)
b_ra_orb = radec_orbitize[:, 0]
b_dec_orb = radec_orbitize[:, 1]
# now project the orbit with rebound
b_manom = basis.tau_to_manom(epochs[0], params[0], params[-1] + params[-3],
params[5], tau_ref_epoch)
c_manom = basis.tau_to_manom(epochs[0], params[0 + 6],
params[-1] + params[-2], params[5 + 6],
tau_ref_epoch)
sim = rebound.Simulation()
sim.units = ('yr', 'AU', 'Msun')
# add star
sim.add(m=params[-1])
# add two planets
sim.add(m=mass_c,
a=params[0 + 6],
e=params[1 + 6],
M=c_manom,
omega=params[3 + 6],
Omega=params[4 + 6] + np.pi / 2,
inc=params[2 + 6])
sim.add(m=mass_b,
a=params[0],
e=params[1],
M=b_manom,
omega=params[3],
Omega=params[4] + np.pi / 2,
inc=params[2])
ps = sim.particles
sim.move_to_com()
# Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period
sim.integrator = "ias15"
sim.dt = ps[1].P / 1000.
# integrate and measure star/planet separation
b_ra_reb = []
b_dec_reb = []
for t in epochs:
sim.integrate(t / 365.25)
b_ra_reb.append(-(ps[2].x - ps[0].x)) # ra is negative x
b_dec_reb.append(ps[2].y - ps[0].y)
b_ra_reb = np.array(b_ra_reb)
b_dec_reb = np.array(b_dec_reb)
diff_ra = b_ra_reb - b_ra_orb / params[6 * 2]
diff_dec = b_dec_reb - b_dec_orb / params[6 * 2]
# should be as good as the one planet case
assert np.all(np.abs(diff_ra) / (params[0] * params[6 * 2]) < 1e-9)
assert np.all(np.abs(diff_dec) / (params[0] * params[6 * 2]) < 1e-9)
def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False):
"""
Calls orbitize.kepler.calc_orbit and optionally accounts for multi-body
interactions. Also computes total quantities like RV (without jitter/gamma)
Args:
params_arr (np.array of float): RxM array
of fitting parameters, where R is the number of
parameters being fit, and M is the number of orbits
we need model predictions for. Must be in the same order
documented in ``System()`` above. If M=1, this can be a 1d array.
epochs (np.array of float): epochs (in mjd) at which to compute
orbit predictions.
comp_rebound (bool, optional): A secondary optional input for
use of N-body solver Rebound; by default, this will be set
to false and a Kepler solver will be used instead.
Returns:
tuple:
raoff (np.array of float): N_epochs x N_bodies x N_orbits array of
RA offsets from barycenter at each epoch.
decoff (np.array of float): N_epochs x N_bodies x N_orbits array of
Dec offsets from barycenter at each epoch.
vz (np.array of float): N_epochs x N_bodies x N_orbits array of
radial velocities at each epoch.
"""
if epochs is None:
epochs = self.data_table['epoch']
n_epochs = len(epochs)
if len(params_arr.shape) == 1:
n_orbits = 1
else:
n_orbits = params_arr.shape[1]
ra_kepler = np.zeros((n_epochs, self.num_secondary_bodies + 1,
n_orbits)) # N_epochs x N_bodies x N_orbits
dec_kepler = np.zeros(
(n_epochs, self.num_secondary_bodies + 1, n_orbits))
ra_perturb = np.zeros(
(n_epochs, self.num_secondary_bodies + 1, n_orbits))
dec_perturb = np.zeros(
(n_epochs, self.num_secondary_bodies + 1, n_orbits))
vz = np.zeros((n_epochs, self.num_secondary_bodies + 1, n_orbits))
# mass/mtot used to compute each Keplerian orbit will be needed later to compute perturbations
if self.track_planet_perturbs:
masses = np.zeros((self.num_secondary_bodies + 1, n_orbits))
mtots = np.zeros((self.num_secondary_bodies + 1, n_orbits))
if comp_rebound or self.use_rebound:
sma = params_arr[self.sma_indx]
ecc = params_arr[self.ecc_indx]
inc = params_arr[self.inc_indx]
argp = params_arr[self.aop_indx]
lan = params_arr[self.pan_indx]
tau = params_arr[self.tau_indx]
plx = params_arr[self.basis.standard_basis_idx['plx']]
if self.fit_secondary_mass:
m_pl = params_arr[self.mpl_idx]
m0 = params_arr[self.basis.param_idx['m0']]
mtot = m0 + sum(m_pl)
else:
m_pl = np.zeros(self.num_secondary_bodies)
# if not fitting for secondary mass, then total mass must be stellar mass
mtot = params_arr[self.basis.param_idx['mtot']]
raoff, deoff, vz = nbody.calc_orbit(
epochs,
sma,
ecc,
inc,
argp,
lan,
tau,
plx,
mtot,
tau_ref_epoch=self.tau_ref_epoch,
m_pl=m_pl,
output_star=True)
else:
for body_num in np.arange(self.num_secondary_bodies) + 1:
sma = params_arr[self.basis.standard_basis_idx['sma{}'.format(
body_num)]]
ecc = params_arr[self.basis.standard_basis_idx['ecc{}'.format(
body_num)]]
inc = params_arr[self.basis.standard_basis_idx['inc{}'.format(
body_num)]]
argp = params_arr[self.basis.standard_basis_idx['aop{}'.format(
body_num)]]
lan = params_arr[self.basis.standard_basis_idx['pan{}'.format(
body_num)]]
tau = params_arr[self.basis.standard_basis_idx['tau{}'.format(
body_num)]]
plx = params_arr[self.basis.standard_basis_idx['plx']]
if self.fit_secondary_mass:
# mass of secondary bodies are in order from -1-num_bodies until -2 in order.
mass = params_arr[self.basis.standard_basis_idx[
'm{}'.format(body_num)]]
m0 = params_arr[self.basis.standard_basis_idx['m0']]
# For what mtot to use to calculate central potential, we should use the mass enclosed in a sphere with r <= distance of planet.
# We need to select all planets with sma < this planet.
all_smas = params_arr[self.sma_indx]
within_orbit = np.where(all_smas <= sma)
outside_orbit = np.where(all_smas > sma)
all_pl_masses = params_arr[self.secondary_mass_indx]
inside_masses = all_pl_masses[within_orbit]
mtot = np.sum(inside_masses) + m0
else:
m_pl = np.zeros(self.num_secondary_bodies)
# if not fitting for secondary mass, then total mass must be stellar mass
mass = None
m0 = None
mtot = params_arr[self.basis.standard_basis_idx['mtot']]
if self.track_planet_perturbs:
masses[body_num] = mass
mtots[body_num] = mtot
# solve Kepler's equation
raoff, decoff, vz_i = kepler.calc_orbit(
epochs,
sma,
ecc,
inc,
argp,
lan,
tau,
plx,
mtot,
mass_for_Kamp=m0,
tau_ref_epoch=self.tau_ref_epoch)
# raoff, decoff, vz are scalers if the length of epochs is 1
if len(epochs) == 1:
raoff = np.array([raoff])
decoff = np.array([decoff])
vz_i = np.array([vz_i])
# add Keplerian ra/deoff for this body to storage arrays
ra_kepler[:, body_num, :] = np.reshape(raoff,
(n_epochs, n_orbits))
dec_kepler[:,
body_num, :] = np.reshape(decoff,
(n_epochs, n_orbits))
vz[:, body_num, :] = np.reshape(vz_i, (n_epochs, n_orbits))
# vz_i is the ith companion radial velocity
if self.fit_secondary_mass:
vz0 = np.reshape(
vz_i * -(mass / m0), (n_epochs, n_orbits)
) # calculating stellar velocity due to ith companion
vz[:, 0, :] += vz0 # adding stellar velocity and gamma
# if we are fitting for the mass of the planets, then they will perturb the star
# add the perturbation on the star due to this planet on the relative astrometry of the planet that was measured
# We are superimposing the Keplerian orbits, so we can add it linearly, scaled by the mass.
# Because we are in Jacobi coordinates, for companions, we only should model the effect of planets interior to it.
# (Jacobi coordinates mean that separation for a given companion is measured relative to the barycenter of all interior companions)
if self.track_planet_perturbs:
for body_num in np.arange(self.num_secondary_bodies + 1):
if body_num > 0:
# for companions, only perturb companion orbits at larger SMAs than this one.
sma = params_arr[self.basis.standard_basis_idx[
'sma{}'.format(body_num)]]
all_smas = params_arr[self.sma_indx]
outside_orbit = np.where(all_smas > sma)[0]
which_perturb_bodies = outside_orbit + 1
# the planet will also perturb the star
which_perturb_bodies = np.append([0],
which_perturb_bodies)
else:
# for the star, what we are measuring is its position relative to the system barycenter
# so we want to account for all of the bodies.
which_perturb_bodies = np.arange(
self.num_secondary_bodies + 1)
for other_body_num in which_perturb_bodies:
# skip itself since the the 2-body problem is measuring the planet-star separation already
if (body_num == other_body_num) | (body_num == 0):
continue
## NOTE: we are only handling astrometry right now (TODO: integrate RV into this)
# this computes the perturbation on the other body due to the current body
# star is perturbed in opposite direction
if other_body_num == 0:
ra_perturb[:, other_body_num, :] -= (
masses[body_num] /
mtots[body_num]) * ra_kepler[:, body_num, :]
dec_perturb[:, other_body_num, :] -= (
masses[body_num] /
mtots[body_num]) * dec_kepler[:, body_num, :]
else:
ra_perturb[:, other_body_num, :] += (
masses[body_num] /
mtots[body_num]) * ra_kepler[:, body_num, :]
dec_perturb[:, other_body_num, :] += (
masses[body_num] /
mtots[body_num]) * dec_kepler[:, body_num, :]
raoff = ra_kepler + ra_perturb
deoff = dec_kepler + dec_perturb
if self.fitting_basis == 'XYZ':
# Find and filter out unbound orbits
bad_orbits = np.where(np.logical_or(ecc >= 1., ecc < 0.))[0]
if (bad_orbits.size != 0):
raoff[:, :, bad_orbits] = np.inf
deoff[:, :, bad_orbits] = np.inf
vz[:, :, bad_orbits] = np.inf
return raoff, deoff, vz
else:
return raoff, deoff, vz
else:
return raoff, deoff, vz
def compute_sep(df,
epochs,
basis,
m0,
m0_err,
plx,
plx_err,
n_planets=1,
pl_num=1):
"""
Computes a sky-projected angular separation posterior given a
RadVel-computed DataFrame.
Args:
df (pd.DataFrame): Radvel-computed posterior (in any orbital basis)
epochs (np.array of astropy.time.Time): epochs at which to compute
separations
basis (str): basis string of input posterior (see
radvel.basis.BASIS_NAMES` for the full list of possibilities).
m0 (float): median of primary mass distribution (assumed Gaussian).
m0_err (float): 1sigma error of primary mass distribution
(assumed Gaussian).
plx (float): median of parallax distribution (assumed Gaussian).
plx_err: 1sigma error of parallax distribution (assumed Gaussian).
n_planets (int): total number of planets in RadVel posterior
pl_num (int): planet number used in RadVel fits (e.g. a RadVel label of
'per1' implies `pl_num` == 1)
Example:
>> df = pandas.read_csv('sample_radvel_chains.csv.bz2', index_col=0)
>> epochs = astropy.time.Time([2022, 2024], format='decimalyear')
>> seps, df_orb = compute_sep(
df, epochs, 'per tc secosw sesinw k', 0.82, 0.02, 312.22, 0.47
)
Returns:
tuple of:
np.array of size (len(epochs) x len(df)): sky-projected angular
separations [mas] at each input epoch
pd.DataFrame: corresponding orbital posterior in orbitize basis
"""
myBasis = Basis(basis, n_planets)
df = myBasis.to_synth(df)
chain_len = len(df)
tau_ref_epoch = 58849
# convert RadVel posteriors -> orbitize posteriors
m_st = np.random.normal(m0, m0_err, size=chain_len)
semiamp = df['k{}'.format(pl_num)].values
per_day = df['per{}'.format(pl_num)].values
period_yr = per_day / 365.25
ecc = df['e{}'.format(pl_num)].values
msini = (Msini(semiamp, per_day, m_st, ecc, Msini_units='Earth') *
(u.M_earth / u.M_sun).to(''))
cosi = (2. * np.random.random(size=chain_len)) - 1.
inc = np.arccos(cosi)
m_pl = msini / np.sin(inc)
mtot = m_st + m_pl
sma = (period_yr**2 * mtot)**(1 / 3)
omega_st_rad = df['w{}'.format(pl_num)].values
omega_pl_rad = omega_st_rad + np.pi
parallax = np.random.normal(plx, plx_err, size=chain_len)
lan = np.random.random_sample(size=chain_len) * 2. * np.pi
tp_mjd = df['tp{}'.format(pl_num)].values - 2400000.5
tau = tp_to_tau(tp_mjd, tau_ref_epoch, period_yr)
# compute projected separation in mas
raoff, deoff, _ = calc_orbit(epochs.mjd,
sma,
ecc,
inc,
omega_pl_rad,
lan,
tau,
parallax,
mtot,
tau_ref_epoch=tau_ref_epoch)
seps = np.sqrt(raoff**2 + deoff**2)
df_orb = pd.DataFrame(np.transpose(
[sma, ecc, inc, omega_pl_rad, lan, tau, parallax, m_st, m_pl]),
columns=[
'sma', 'ecc', 'inc_rad', 'omega_pl_rad',
'lan_rad', 'tau_58849', 'plx', 'm_st', 'mp'
])
return seps, df_orb
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='epochs', mod180=False):
"""
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: epochs, sma1, ecc1, inc1, aop1,
pan1, tau1. Number can be switched out. Default is epochs.
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 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 Exception('start_mjd keyword date must be less than sep_pa_end_year keyword date.')
with warnings.catch_warnings():
warnings.simplefilter('ignore', ErfaWarning)
dict_of_indices = {
'sma': 0,
'ecc': 1,
'inc': 2,
'aop': 3,
'pan': 4,
'tau': 5,
'plx':6
}
if cbar_param == 'epochs':
pass
elif cbar_param[0:3] in dict_of_indices:
try:
object_id = np.int(cbar_param[3:])
except ValueError:
object_id = 1
index = dict_of_indices[cbar_param[0:3]] + 6*(object_id-1)
else:
raise Exception('Invalid input; acceptable inputs include epochs, sma1, ecc1, inc1, aop1, pan1, tau1, sma2, ecc2, ...')
# Split the 2-D post array into series of 1-D arrays for each orbital parameter
num_objects, remainder = np.divmod(self.post.shape[1],6)
if object_to_plot > num_objects:
return None
sma = self.post[:,dict_of_indices['sma']]
ecc = self.post[:,dict_of_indices['ecc']]
inc = self.post[:,dict_of_indices['inc']]
aop = self.post[:,dict_of_indices['aop']]
pan = self.post[:,dict_of_indices['pan']]
tau = self.post[:,dict_of_indices['tau']]
plx = self.post[:,dict_of_indices['plx']]
# Then, get the other parameters
if 'mtot' in self.labels:
mtot = self.post[:,-1]
elif 'm0' in self.labels:
m0 = self.post[:,-1]
mplanet = self.post[:,-2]
mtot = m0 + mplanet
# Select random indices for plotted orbit
if num_orbits_to_plot > len(sma):
num_orbits_to_plot = len(sma)
choose = np.random.randint(0, high=len(sma), size=num_orbits_to_plot)
raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
deoff = 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):
orb_ind = choose[i]
# 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[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],
tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch
)
raoff[i,:] = raoff0
deoff[i,:] = deoff0
# Create a linearly increasing colormap for our range of epochs
if cbar_param != 'epochs':
cbar_param_arr = self.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 == 'epochs':
norm = mpl.colors.Normalize(vmin=np.min(epochs), vmax=np.max(epochs[-1,:]))
norm_yr = mpl.colors.Normalize(
vmin=np.min(Time(epochs,format='mjd').decimalyear),
vmax=np.max(Time(epochs,format='mjd').decimalyear)
)
# Create figure for orbit plots
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):
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 != 'epochs':
lc.set_array(np.ones(len(epochs[0]))*cbar_param_arr[i])
elif cbar_param == 'epochs':
lc.set_array(epochs[i,:])
ax.add_collection(lc)
# 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
# add colorbar
if show_colorbar:
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', label=cbar_param)
# 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], tau_ref_epoch=self.tau_ref_epoch
)
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,:], mod180=mod180)
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
def compute_model(self, params_arr):
"""
Compute model predictions for an array of fitting parameters.
Args:
params_arr (np.array of float): RxM array
of fitting parameters, where R is the number of
parameters being fit, and M is the number of orbits
we need model predictions for. Must be in the same order
documented in ``System()`` above. If M=1, this can be a 1d array.
Returns:
np.array of float: Nobsx2xM array model predictions. If M=1, this is
a 2d array, otherwise it is a 3d array.
"""
if len(params_arr.shape) == 1:
model = np.zeros((len(self.data_table), 2))
# Rob and Lea: adding gamma zeros
gamma = np.zeros((len(self.data_table), 2))
jitter = np.zeros((len(self.data_table), 2))
else:
model = np.zeros((len(self.data_table), 2, params_arr.shape[1]))
gamma = np.zeros((len(self.data_table), 2, params_arr.shape[1]))
jitter = np.zeros((len(self.data_table), 2, params_arr.shape[1]))
if self.track_planet_perturbs:
radec_perturb = np.zeros(model.shape)
if len(self.rv[0]) > 0 and self.fit_secondary_mass:
# looping through instruments to get the gammas
for rv_idx in range(len(self.rv_instruments)):
gamma[self.rv_inst_indices[rv_idx],
0] = params_arr[6 * self.num_secondary_bodies + 1 +
2 * rv_idx] # km/s
jitter[self.rv_inst_indices[rv_idx],
0] = params_arr[6 * self.num_secondary_bodies + 2 +
2 * rv_idx]
gamma[self.rv_inst_indices[rv_idx], 1] = np.nan
jitter[self.rv_inst_indices[rv_idx], 1] = np.nan
# need to put planetary rv later
total_rv0 = 0 # If we're not fitting rv, then we don't regard the total rv and will not use this
for body_num in np.arange(self.num_secondary_bodies) + 1:
# we're going to compute at all epochs for convenience of indexing right now
# self.radec, and self.seppa index into the entire data table, not just the values for a particular body
epochs = self.data_table['epoch'] #[self.body_indices[body_num]]
startindex = 6 * (body_num - 1)
sma = params_arr[startindex]
ecc = params_arr[startindex + 1]
inc = params_arr[startindex + 2]
argp = params_arr[startindex + 3]
lan = params_arr[startindex + 4]
tau = params_arr[startindex + 5]
plx = params_arr[6 * self.num_secondary_bodies]
if self.fit_secondary_mass:
# mass of secondary bodies are in order from -1-num_bodies until -2 in order.
mass = params_arr[-1 - self.num_secondary_bodies +
(body_num - 1)]
m0 = params_arr[-1]
# For what mtot to use to calculate central potential, we should use the mass enclosed in a sphere with r <= distance of planet.
# We need to select all planets with sma < this planet.
all_smas = params_arr[0:6 * self.num_secondary_bodies:6]
within_orbit = np.where(all_smas <= sma)
outside_orbit = np.where(all_smas > sma)
all_pl_masses = params_arr[-1 - self.num_secondary_bodies:-1]
inside_masses = all_pl_masses[within_orbit]
mtot = np.sum(inside_masses) + m0
else:
# if not fitting for secondary mass, then total mass must be stellar mass
mass = None
m0 = None
mtot = params_arr[-1]
# i = 1,2,3... (companion index)
raoff, decoff, vz_i = kepler.calc_orbit(
epochs,
sma,
ecc,
inc,
argp,
lan,
tau,
plx,
mtot,
mass_for_Kamp=m0,
tau_ref_epoch=self.tau_ref_epoch,
tau_warning=False)
# raoff, decoff, vz are scalers if the length of epochs is 1.
# Jason is too lazy to figure out how to make it return a one element array without breaking everything else
# so hard code it here to convert them into 1-element numpy arrays.
if len(epochs) == 1:
raoff = np.array([raoff])
decoff = np.array([decoff])
vz = np.array([vz_i])
# vz_i is the ith companion radial velocity
if len(self.rv[0]) > 0 and self.fit_secondary_mass:
vz0 = vz_i * -(
mass / m0
) # calculating stellar velocity due to ith companion
total_rv0 = total_rv0 + vz0 # Adding stellar velocity and gamma
# for the model points that correspond to this planet's orbit, add the model prediction
# RA/Dec
if len(self.radec[body_num]
) > 0: # (prevent empty array dimension errors)
model[self.radec[body_num], 0] = raoff[self.radec[body_num]]
model[self.radec[body_num], 1] = decoff[self.radec[body_num]]
# Sep/PA
if len(self.seppa[body_num]) > 0:
sep, pa = radec2seppa(raoff, decoff)
model[self.seppa[body_num], 0] = sep[self.seppa[body_num]]
model[self.seppa[body_num], 1] = pa[self.seppa[body_num]]
# RV
if len(self.rv[body_num]) > 0:
model[self.rv[body_num], 0] = vz_i[self.rv[body_num]]
model[self.rv[body_num], 1] = np.nan
# for the other epochs, if we are fitting for the mass of the planets, then they will perturb the star
# add the perturbation on the star due to this planet on the relative astrometry of the planet that was measured
# We are superimposing the Keplerian orbits, so we can add it linearly, scaled by the mass.
# Because we are in Jacobi coordinates, for companions, we only should model the effect of planets interior to it.
# (Jacobi coordinates mean that separation for a given companion is measured relative to the barycenter of all interior companions)
if self.track_planet_perturbs:
if body_num > 0:
# for companions, only perturb companion orbits at larger SMAs than this one.
# note the +1, since the 0th planet is body_num 1.
which_perturb_bodies = outside_orbit[0] + 1
else:
# for the star, what we are measuring is it's position relative to the system barycenter
# so we want to account for all of the bodies.
which_perturb_bodies = range(self.num_secondary_bodies + 1)
for other_body_num in which_perturb_bodies:
# skip itself since the the 2-body problem is measuring the planet-star separation already
if (body_num == other_body_num) | (body_num == 0):
continue
## NOTE: we are only handling ra/dec and sep/pa right now
## TODO: integrate RV into this
if len(self.radec[other_body_num]) > 0:
radec_perturb[self.radec[other_body_num],
0] += -(mass / mtot) * raoff[
self.radec[other_body_num]]
radec_perturb[self.radec[other_body_num],
1] += -(mass / mtot) * decoff[
self.radec[other_body_num]]
if len(self.seppa[other_body_num]) > 0:
radec_perturb[self.seppa[other_body_num],
0] += -(mass / mtot) * raoff[
self.seppa[other_body_num]]
radec_perturb[self.seppa[other_body_num],
1] += -(mass / mtot) * decoff[
self.seppa[other_body_num]]
if len(self.rv[0]) > 0 and self.fit_secondary_mass:
if len(total_rv0[self.rv[0]]) > 0:
model[self.rv[0], 0] = total_rv0[self.rv[0]]
model[self.rv[0], 1] = np.nan # nans only for rv indices
# add the effects of other planets on the measured astrometry
if self.track_planet_perturbs:
for body_num in range(self.num_secondary_bodies + 1):
if len(self.radec[body_num]) > 0:
model[self.radec[body_num]] -= radec_perturb[
self.radec[body_num]]
if len(self.seppa[body_num]) > 0:
# for seppa, add the perturbations in radec space and convert back
ra_unperturb, dec_unperturb = seppa2radec(
model[self.seppa[body_num], 0],
model[self.seppa[body_num], 1])
ra_perturb = ra_unperturb - radec_perturb[
self.seppa[body_num], 0]
dec_perturb = dec_unperturb - radec_perturb[
self.seppa[body_num], 1]
sep_perturb, pa_perturb = radec2seppa(
ra_perturb, dec_perturb)
model[self.seppa[body_num], 0] = sep_perturb
model[self.seppa[body_num], 1] = pa_perturb
if self.fit_secondary_mass:
return model + gamma, jitter
else:
return model, jitter
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, 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.
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.")
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')
sep_data, sep_err=data['quant1'][seppa_inds],data['quant1_err'][seppa_inds]
pa_data, pa_err=data['quant2'][seppa_inds],data['quant2_err'][seppa_inds]
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 = np.append(sep_data, sep_from_ra_data)
sep_err = np.append(sep_err, sep_err_from_ra_data)
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)
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:
ra_data,dec_data=orbitize.system.seppa2radec(sep_data,pa_data)
# 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'
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)
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:
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]
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]
med_ga=[results.post[best_like,i] for i in gam_idx]
# 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=('#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-=med_ga[i]
rvs -= best_post[results.param_idx[gams[i]]]
plt.scatter(epochs,rvs,s=5,marker=next(ax3_symbols),c=next(ax3_colors),label=name,zorder=5)
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)
# 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
"""
A script to test some edge cases for the calc_orbit() function. We threw a bunch of weird inputs in
and now everything is returning nans. Try to figure out which input or inputs is causing it to return nan.
If you have time, update the code to do some error checking to alert the user the input is invalid.
"""
import numpy as np
import orbitize.kepler as kepler
ra, dec, rv = kepler.calc_orbit(np.array([1000, 1e6, -12]), 10.0, 0.1, -1,
1000, 0, 0.5, 100, -99)
print(ra, dec, rv)
assert np.all(np.isfinite(ra))
assert np.all(np.isfinite(dec))
assert np.all(np.isfinite(rv))
def test_fit_selfconsist():
"""
Tests that the masses we get from orbitize! are what we expeect. Note that this does not test for correctness.
"""
# generate planet b orbital parameters
b_params = [1, 0, 0, 0, 0, 0.5]
tau_ref_epoch = 0
mass_b = 0.001 # Msun
m0 = 1 # Msun
plx = 1 # mas
# generate planet c orbital parameters
# at 0.3 au, and starts on the opposite side of the star relative to b
c_params = [0.3, 0, 0, np.pi, 0, 0.5]
mass_c = 0.002 # Msun
mtot_c = m0 + mass_b + mass_c
mtot_b = m0 + mass_b
period_b = np.sqrt(b_params[0]**3 / mtot_b)
period_c = np.sqrt(c_params[0]**3 / mtot_c)
epochs = np.linspace(0, period_b * 365.25,
20) + tau_ref_epoch # the full period of b, MJD
# comptue Keplerian orbit of b
ra_model_b, dec_model_b, vz_model = kepler.calc_orbit(
epochs,
b_params[0],
b_params[1],
b_params[2],
b_params[3],
b_params[4],
b_params[5],
plx,
mtot_b,
mass_for_Kamp=m0,
tau_ref_epoch=tau_ref_epoch)
# comptue Keplerian orbit of c
ra_model_c, dec_model_c, vz_model_c = kepler.calc_orbit(
epochs,
c_params[0],
c_params[1],
c_params[2],
c_params[3],
c_params[4],
c_params[5],
plx,
mtot_c,
tau_ref_epoch=tau_ref_epoch)
# perturb b due to c
ra_model_b_orig = np.copy(ra_model_b)
dec_model_b_orig = np.copy(dec_model_b)
# the sign is positive b/c of 2 negative signs cancelling.
ra_model_b += mass_c / m0 * ra_model_c
dec_model_b += mass_c / m0 * dec_model_c
# # perturb c due to b
# ra_model_c += mass_b/m0 * ra_model_b_orig
# dec_model_c += mass_b/m0 * dec_model_b_orig
# generate some fake measurements to fit to. Make it with b first
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int), ra_model_b,
0.00001 * np.ones(epochs.shape, dtype=int), dec_model_b,
0.00001 * np.ones(epochs.shape, dtype=int)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
# add c
for eps, ra, dec in zip(epochs, ra_model_c, dec_model_c):
t.add_row([eps, 2, ra, 0.000001, dec, 0.000001])
filename = os.path.join(orbitize.DATADIR,
"multiplanet_fake_2planettest.csv")
t.write(filename, overwrite=True)
# create the orbitize system and generate model predictions using the standard 1 body model for b, and the 2 body model for b and c
astrom_dat = read_input.read_file(filename)
sys = system.System(2,
astrom_dat,
m0,
plx,
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
# fix most of the orbital parameters to make the dimensionality a bit smaller
sys.sys_priors[sys.param_idx['ecc1']] = b_params[1]
sys.sys_priors[sys.param_idx['inc1']] = b_params[2]
sys.sys_priors[sys.param_idx['aop1']] = b_params[3]
sys.sys_priors[sys.param_idx['pan1']] = b_params[4]
sys.sys_priors[sys.param_idx['ecc2']] = c_params[1]
sys.sys_priors[sys.param_idx['inc2']] = c_params[2]
sys.sys_priors[sys.param_idx['aop2']] = c_params[3]
sys.sys_priors[sys.param_idx['pan2']] = c_params[4]
sys.sys_priors[sys.param_idx['m1']] = priors.LogUniformPrior(
mass_b * 0.01, mass_b * 100)
sys.sys_priors[sys.param_idx['m2']] = priors.LogUniformPrior(
mass_c * 0.01, mass_c * 100)
n_walkers = 30
samp = sampler.MCMC(sys, num_temps=1, num_walkers=n_walkers, num_threads=1)
# should have 8 parameters
assert samp.num_params == 6
# start walkers near the true location for the orbital parameters
np.random.seed(123)
# planet b
samp.curr_pos[:, 0] = np.random.normal(b_params[0], 0.01, n_walkers) # sma
samp.curr_pos[:, 1] = np.random.normal(b_params[-1], 0.01,
n_walkers) # tau
# planet c
samp.curr_pos[:, 2] = np.random.normal(c_params[0], 0.01, n_walkers) # sma
samp.curr_pos[:, 3] = np.random.normal(c_params[-1], 0.01,
n_walkers) # tau
# we will make a fairly broad mass starting position
samp.curr_pos[:, 4] = np.random.uniform(mass_b * 0.25, mass_b * 4,
n_walkers)
samp.curr_pos[:, 5] = np.random.uniform(mass_c * 0.25, mass_c * 4,
n_walkers)
samp.curr_pos[0, 4] = mass_b
samp.curr_pos[0, 5] = mass_c
samp.run_sampler(n_walkers * 50, burn_steps=600)
res = samp.results
print(np.median(res.post[:, sys.param_idx['m1']]),
np.median(res.post[:, sys.param_idx['m2']]))
assert np.median(res.post[:, sys.param_idx['sma1']]) == pytest.approx(
b_params[0], abs=0.01)
assert np.median(res.post[:, sys.param_idx['sma2']]) == pytest.approx(
c_params[0], abs=0.01)
assert np.median(res.post[:, sys.param_idx['m2']]) == pytest.approx(
mass_c, abs=0.5 * mass_c)
def test_1planet():
"""
Sanity check that things agree for 1 planet case
"""
# generate a planet orbit
sma = 1
ecc = 0.1
inc = np.radians(45)
aop = np.radians(45)
pan = np.radians(45)
tau = 0.5
plx = 1
mtot = 1
tau_ref_epoch = 0
mjup = u.Mjup.to(u.Msun)
mass_b = 12 * mjup
epochs = np.linspace(0, 300,
100) + tau_ref_epoch # nearly the full period, MJD
ra_model, dec_model, vz_model = kepler.calc_orbit(
epochs,
sma,
ecc,
inc,
aop,
pan,
tau,
plx,
mtot,
tau_ref_epoch=tau_ref_epoch)
# generate some fake measurements just to feed into system.py to test bookkeeping
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int), ra_model,
np.zeros(ra_model.shape), dec_model,
np.zeros(dec_model.shape)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
filename = os.path.join(orbitize.DATADIR, "rebound_1planet.csv")
t.write(filename)
# create the orbitize system and generate model predictions using the ground truth
astrom_dat = read_input.read_file(filename)
sys = system.System(1, astrom_dat, mtot, plx, tau_ref_epoch=tau_ref_epoch)
params = np.array([sma, ecc, inc, aop, pan, tau, plx, mtot])
radec_orbitize, _ = sys.compute_model(params)
ra_orb = radec_orbitize[:, 0]
dec_orb = radec_orbitize[:, 1]
# now project the orbit with rebound
manom = basis.tau_to_manom(epochs[0], sma, mtot, tau, tau_ref_epoch)
sim = rebound.Simulation()
sim.units = ('yr', 'AU', 'Msun')
# add star
sim.add(m=mtot - mass_b)
# add one planet
sim.add(m=mass_b,
a=sma,
e=ecc,
M=manom,
omega=aop,
Omega=pan + np.pi / 2,
inc=inc)
ps = sim.particles
sim.move_to_com()
# Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period
sim.integrator = "ias15"
sim.dt = ps[1].P / 1000.
# integrate and measure star/planet separation
ra_reb = []
dec_reb = []
for t in epochs:
sim.integrate(t / 365.25)
ra_reb.append(-(ps[1].x - ps[0].x)) # ra is negative x
dec_reb.append(ps[1].y - ps[0].y)
ra_reb = np.array(ra_reb)
dec_reb = np.array(dec_reb)
diff_ra = ra_reb - ra_orb / plx
diff_dec = dec_reb - dec_orb / plx
assert np.all(np.abs(diff_ra) < 1e-9)
assert np.all(np.abs(diff_dec) < 1e-9)
def test_2planet_massive():
"""
Compare multiplanet to rebound for planets with mass.
"""
# generate a planet orbit
mjup = u.Mjup.to(u.Msun)
mass_b = 12 * mjup
mass_c = 9 * mjup
params = np.array([
10, 0.1,
np.radians(45),
np.radians(45),
np.radians(45), 0.5, 3, 0.1,
np.radians(45),
np.radians(190),
np.radians(45), 0.2, 50, mass_b, mass_c, 1.5 - mass_b - mass_c
])
params_noc = np.array([
10, 0.1,
np.radians(45),
np.radians(45),
np.radians(45), 0.5, 3, 0.1,
np.radians(45),
np.radians(190),
np.radians(45), 0.2, 50, mass_b, 0, 1.5 - mass_b
])
tau_ref_epoch = 0
epochs = np.linspace(0, 365.25 * 10,
100) + tau_ref_epoch # nearly the full period, MJD
# doesn't matter that this is right, just needs to be the same size. below doesn't include effect of c
# just want to generate some measurements of plaent b to test compute model
b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit(
epochs,
params[0],
params[1],
params[2],
params[3],
params[4],
params[5],
params[-2],
params[-1],
tau_ref_epoch=tau_ref_epoch)
# generate some fake measurements of planet b, just to feed into system.py to test bookkeeping
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int), b_ra_model,
np.zeros(b_ra_model.shape), b_dec_model,
np.zeros(b_dec_model.shape)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
filename = os.path.join(orbitize.DATADIR, "rebound_2planet_outer.csv")
t.write(filename)
#### TEST THE OUTER PLANET ####
# create the orbitize system and generate model predictions using the ground truth
astrom_dat = read_input.read_file(filename)
sys = system.System(2,
astrom_dat,
params[-1],
params[-4],
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
# generate measurement
radec_orbitize, _ = sys.compute_model(params)
b_ra_orb = radec_orbitize[:, 0]
b_dec_orb = radec_orbitize[:, 1]
# debug, generate measurement without c having any mass
radec_orb_noc, _ = sys.compute_model(params_noc)
b_ra_orb_noc = radec_orb_noc[:, 0]
b_dec_orb_noc = radec_orb_noc[:, 1]
# check that planet c's perturbation is imprinted (nonzero))
#assert np.all(b_ra_orb_noc != b_ra_orb)
# now project the orbit with rebound
b_manom = basis.tau_to_manom(epochs[0], params[0], params[-1] + params[-3],
params[5], tau_ref_epoch)
c_manom = basis.tau_to_manom(epochs[0], params[0 + 6],
params[-1] + params[-2], params[5 + 6],
tau_ref_epoch)
sim = rebound.Simulation()
sim.units = ('yr', 'AU', 'Msun')
# add star
sim.add(m=params[-1])
# add two planets
sim.add(m=mass_c,
a=params[0 + 6],
e=params[1 + 6],
M=c_manom,
omega=params[3 + 6],
Omega=params[4 + 6] + np.pi / 2,
inc=params[2 + 6])
sim.add(m=mass_b,
a=params[0],
e=params[1],
M=b_manom,
omega=params[3],
Omega=params[4] + np.pi / 2,
inc=params[2])
ps = sim.particles
sim.move_to_com()
# Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period
sim.integrator = "ias15"
sim.dt = ps[1].P / 1000.
# integrate and measure star/planet separation
b_ra_reb = []
b_dec_reb = []
for t in epochs:
sim.integrate(t / 365.25)
b_ra_reb.append(-(ps[2].x - ps[0].x)) # ra is negative x
b_dec_reb.append(ps[2].y - ps[0].y)
b_ra_reb = np.array(b_ra_reb)
b_dec_reb = np.array(b_dec_reb)
diff_ra = b_ra_reb - b_ra_orb / params[6 * 2]
diff_dec = b_dec_reb - b_dec_orb / params[6 * 2]
# we placed the planets far apart to minimize secular interactions but there are still some, so relax precision
assert np.all(np.abs(diff_ra) / (params[0]) < 1e-3)
assert np.all(np.abs(diff_dec) / (params[0]) < 1e-3)
###### NOW TEST THE INNER PLANET #######
# generate some fake measurements of planet c, just to feed into system.py to test bookkeeping
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int) * 2, b_ra_model,
np.zeros(b_ra_model.shape), b_dec_model,
np.zeros(b_dec_model.shape)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
filename = os.path.join(orbitize.DATADIR, "rebound_2planet_inner.csv")
t.write(filename)
# create the orbitize system and generate model predictions using the ground truth
astrom_dat = read_input.read_file(filename)
sys = system.System(2,
astrom_dat,
params[-1],
params[-2],
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
# generate measurement
radec_orbitize, _ = sys.compute_model(params)
c_ra_orb = radec_orbitize[:, 0]
c_dec_orb = radec_orbitize[:, 1]
# start the REBOUND sim again
sim = rebound.Simulation()
sim.units = ('yr', 'AU', 'Msun')
# add star
sim.add(m=params[-1])
# add two planets
sim.add(m=mass_c,
a=params[0 + 6],
e=params[1 + 6],
M=c_manom,
omega=params[3 + 6],
Omega=params[4 + 6] + np.pi / 2,
inc=params[2 + 6])
sim.add(m=mass_b,
a=params[0],
e=params[1],
M=b_manom,
omega=params[3],
Omega=params[4] + np.pi / 2,
inc=params[2])
ps = sim.particles
sim.move_to_com()
# Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period
sim.integrator = "ias15"
sim.dt = ps[1].P / 1000.
# integrate and measure star/planet separation
c_ra_reb = []
c_dec_reb = []
for t in epochs:
sim.integrate(t / 365.25)
c_ra_reb.append(-(ps[1].x - ps[0].x)) # ra is negative x
c_dec_reb.append(ps[1].y - ps[0].y)
c_ra_reb = np.array(c_ra_reb)
c_dec_reb = np.array(c_dec_reb)
diff_ra = c_ra_reb - c_ra_orb / params[6 * 2]
diff_dec = c_dec_reb - c_dec_orb / params[6 * 2]
# planet is 3 times closer, so roughly 3 times larger secular errors.
assert np.all(np.abs(diff_ra) / (params[0]) < 3e-3)
assert np.all(np.abs(diff_dec) / (params[0]) < 3e-3)
def print_prediction(planet_name,
date_mjd,
chains,
tau_ref_epoch,
num_samples=None):
"""
Prints out a prediction for the prediction of a planet given a set of posterior draws
Args:
planet_name (str): name of the planet, already checked to be in the list
date_mjd (float): MJD of date for which we want a prediction
chains (np.array): Nx8 array of N orbital elements. Orbital elements are ordered as:
sma, ecc, inc, aop, pan, tau, plx, mtot
tau_ref_epoch (float): MJD for reference epoch of tau (see orbitize for details on tau)
num_samples (int): number of random samples for prediction. If None, will use all of them
Returns:
ra_args (tuple): a two-element tuple of the median RA offset, and stddev of RA offset
dec_args (tuple): a two-element tuple of the median Dec offset, and stddev of Dec offset
sep_args (tuple): a two-element tuple of the median separation offset, and stddev of sep offset
pa_args (tuple): a two-element tuple of the median PA offset, and stddev of PA offset
"""
if num_samples is None:
num_samples = chains.shape[0]
rand_draws = np.arange(num_samples) # don't need to randomize
else:
if num_samples > chains.shape[0]:
print("Requested too many samples. Maximum is {0}.".format(
chains.shape[0]))
return
# randomly draw values
rand_draws = np.random.randint(0, chains.shape[0], num_samples)
rand_orbits = chains[rand_draws]
if planet_name not in multi_dict:
# single Keplerian orbit fits
sma = rand_orbits[:, 0]
ecc = rand_orbits[:, 1]
inc = rand_orbits[:, 2]
aop = rand_orbits[:, 3]
pan = rand_orbits[:, 4]
tau = rand_orbits[:, 5]
plx = rand_orbits[:, 6]
mtot = rand_orbits[:, 7]
rand_ras, rand_decs, rand_vzs = kepler.calc_orbit(
date_mjd,
sma,
ecc,
inc,
aop,
pan,
tau,
plx,
mtot,
tau_ref_epoch=tau_ref_epoch)
else:
# massive multi-planet orbit fits
planet_num, tot_planets = multi_dict[planet_name]
orb_indices = np.arange(6 * planet_num, 6 * planet_num + 6, 1)
sma, ecc, inc, aop, pan, tau = rand_orbits[:, orb_indices].T
plx = rand_orbits[:, 6 * tot_planets]
mass_planets = rand_orbits[:, -1 - tot_planets:-1]
mass_star = rand_orbits[:, -1]
all_pl_smas = rand_orbits[0, 0:6 * tot_planets:6]
within_orbit = np.where(all_pl_smas <= all_pl_smas[planet_num])
mtot = mass_star + np.sum(mass_planets[:, within_orbit[0]], axis=1)
rand_ras, rand_decs, rand_vzs = kepler.calc_orbit(
date_mjd,
sma,
ecc,
inc,
aop,
pan,
tau,
plx,
mtot,
tau_ref_epoch=tau_ref_epoch)
# add perturbation from other planets
for inner_pl in within_orbit[0]:
if inner_pl == planet_num:
continue
within_inner_orbit = np.where(all_pl_smas < all_pl_smas[inner_pl])
inner_mtot = mass_star + np.sum(
mass_planets[:, within_inner_orbit[0]], axis=1)
mass_inner = mass_planets[:, inner_pl]
inner_orb_indices = np.arange(6 * inner_pl, 6 * inner_pl + 6, 1)
in_sma, in_ecc, in_inc, in_aop, in_pan, in_tau = rand_orbits[:,
inner_orb_indices].T
inner_rand_ras, inner_rand_decs, inner_rand_vzs = kepler.calc_orbit(
date_mjd,
in_sma,
in_ecc,
in_inc,
in_aop,
in_pan,
in_tau,
plx,
inner_mtot,
tau_ref_epoch=tau_ref_epoch)
rand_ras += mass_inner / inner_mtot * inner_rand_ras
rand_decs += mass_inner / inner_mtot * inner_rand_decs
rand_vzs += mass_inner / inner_mtot * inner_rand_vzs
rand_seps = np.sqrt(rand_ras**2 + rand_decs**2)
rand_pas = np.degrees(np.arctan2(rand_ras, rand_decs)) % 360
ra_args = np.median(rand_ras), np.std(rand_ras)
dec_args = np.median(rand_decs), np.std(rand_decs)
sep_args = np.median(rand_seps), np.std(rand_seps)
pa_args = np.median(rand_pas), np.std(rand_pas)
rv_args = np.median(rand_vzs), np.std(rand_vzs)
print("RA Offset = {0:.3f} +/- {1:.3f} mas".format(ra_args[0], ra_args[1]))
print("Dec Offset = {0:.3f} +/- {1:.3f} mas".format(
dec_args[0], dec_args[1]))
print("Separation = {0:.3f} +/- {1:.3f} mas".format(
sep_args[0], sep_args[1]))
print("PA = {0:.3f} +/- {1:.3f} deg".format(pa_args[0], pa_args[1]))
print("Planetary RV = {0:.3f} +/- {1:.3f} km/s".format(
rv_args[0], rv_args[1]))
return ra_args, dec_args, sep_args, pa_args
def test_1planet():
"""
Check that for the 2-body case, the primary orbit around the barycenter
is equal to -m2/(m1 + m2) times the secondary orbit around the primary.
"""
# generate a planet orbit
sma = 1
ecc = 0.1
inc = np.radians(45)
aop = np.radians(45)
pan = np.radians(45)
tau = 0.5
plx = 1
m0 = 1
tau_ref_epoch = 0
mjup = u.Mjup.to(u.Msun)
mass_b = 100 * mjup
mtot = mass_b + m0
epochs = np.linspace(0, 300,
100) + tau_ref_epoch # nearly the full period, MJD
ra_model, dec_model, _ = kepler.calc_orbit(epochs,
sma,
ecc,
inc,
aop,
pan,
tau,
plx,
mtot,
tau_ref_epoch=tau_ref_epoch)
# generate some fake measurements to feed into system.py to test bookkeeping
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int), ra_model,
np.zeros(ra_model.shape), dec_model,
np.zeros(dec_model.shape)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
filename = os.path.join(orbitize.DATADIR, "rebound_1planet.csv")
t.write(filename)
# create the orbitize system and generate model predictions using ground truth
astrom_dat = read_input.read_file(filename)
sys = system.System(1,
astrom_dat,
m0,
plx,
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
sys.track_planet_perturbs = True
params = np.array([sma, ecc, inc, aop, pan, tau, plx, mass_b, m0])
ra, dec, _ = sys.compute_all_orbits(params)
# the planet and stellar orbit should just be scaled versions of one another
planet_ra = ra[:, 1, :]
planet_dec = dec[:, 1, :]
star_ra = ra[:, 0, :]
star_dec = dec[:, 0, :]
assert np.all(np.abs(star_ra + (mass_b / mtot) * planet_ra) < 1e-16)
assert np.all(np.abs(star_dec + (mass_b / mtot) * planet_dec) < 1e-16)
# remove the created csv file to clean up
os.system('rm {}'.format(filename))
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):
"""
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: epochs, sma1, ecc1, inc1, aop1,
pan1, tau1. Number can be switched out. Default is epochs.
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.
astrometry (Boolean): set to True by default. Plots the astrometric data.
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 > self.num_secondary_bodies:
raise ValueError(
"Only {0} secondary bodies being fit. Requested to plot body {1} which is out of range"
.format(self.num_secondary_bodies, object_to_plot))
if object_to_plot == 0:
raise ValueError(
"Plotting the primary's orbit is currently unsupported. Stay tuned.."
)
with warnings.catch_warnings():
warnings.simplefilter('ignore', ErfaWarning)
dict_of_indices = {
'sma': 0,
'ecc': 1,
'inc': 2,
'aop': 3,
'pan': 4,
'tau': 5,
'plx': 6 * self.num_secondary_bodies,
}
if cbar_param == 'Epoch [year]':
pass
elif cbar_param[0:3] in dict_of_indices:
try:
object_id = np.int(cbar_param[3:])
except ValueError:
object_id = 1
index = dict_of_indices[cbar_param[0:3]] + 6 * (object_id - 1)
else:
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']]
ecc = self.post[:, start_index + dict_of_indices['ecc']]
inc = self.post[:, start_index + dict_of_indices['inc']]
aop = self.post[:, start_index + dict_of_indices['aop']]
pan = self.post[:, start_index + dict_of_indices['pan']]
tau = self.post[:, start_index + dict_of_indices['tau']]
plx = self.post[:, dict_of_indices['plx']]
# Then, get the other parameters
if 'mtot' in self.labels:
mtot = self.post[:, -1]
elif 'm0' in self.labels:
m0 = self.post[:, -1]
m1 = self.post[:, -(self.num_secondary_bodies + 1) +
(object_to_plot - 1)]
mtot = m0 + m1
# Select random indices for plotted orbit
if num_orbits_to_plot > len(sma):
num_orbits_to_plot = len(sma)
choose = np.random.randint(0,
high=len(sma),
size=num_orbits_to_plot)
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):
orb_ind = choose[i]
# 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[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],
tau[orb_ind],
plx[orb_ind],
mtot[orb_ind],
tau_ref_epoch=self.tau_ref_epoch,
tau_warning=False)
raoff[i, :] = raoff0
deoff[i, :] = deoff0
# Create a linearly increasing colormap for our range of epochs
if cbar_param != 'Epoch [year]':
cbar_param_arr = self.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]':
norm = mpl.colors.Normalize(vmin=np.min(epochs),
vmax=np.max(epochs[-1, :]))
norm_yr = mpl.colors.Normalize(
vmin=np.min(Time(epochs, format='mjd').decimalyear),
vmax=np.max(Time(epochs, format='mjd').decimalyear))
# 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)
data = self.data
astr_inds = np.where((~np.isnan(data['quant1']))
& (~np.isnan(data['quant2'])))
astr_epochs = data['epoch'][astr_inds]
sep_data, sep_err = data['quant1'][astr_inds], data['quant1_err'][
astr_inds]
pa_data, pa_err = data['quant2'][astr_inds], data['quant2_err'][
astr_inds]
# 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:
ra_data, dec_data = orbitize.system.seppa2radec(
sep_data, pa_data)
ax.scatter(ra_data,
dec_data,
marker='*',
c='#FF7F11',
zorder=10,
s=60)
# 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
# Rob: Moved colorbar size to the bottom after tight_layout() because the cbar scaling was not compatible with tight_layout()
# 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')
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
if rv_time_series:
raoff0, deoff0, vzoff0 = 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],
tau_ref_epoch=self.tau_ref_epoch,
mass_for_Kamp=m0[orb_ind],
tau_warning=False)
raoff[i, :] = raoff0
deoff[i, :] = deoff0
else:
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],
tau_ref_epoch=self.tau_ref_epoch,
tau_warning=False)
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)
# 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)
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 gamma/sigma labels and corresponding positions in the posterior
gams = ['gamma_' + inst for inst in insts]
if isinstance(self.labels, list):
labels = np.array(self.labels)
else:
labels = self.labels
# get the indices corresponding to each gamma within self.labels
gam_idx = [
np.where(labels == inst_gamma)[0][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(
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]
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']
# get rvs and plot them
for i, name in enumerate(inds.keys()):
rv_inds = np.where((np.isnan(data['quant2'])))
inst_data = data[inds[name]]
rvs = inst_data['quant1']
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)
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],
tau[best_like],
plx[best_like],
mtot[best_like],
tau_ref_epoch=self.tau_ref_epoch,
mass_for_Kamp=m0[best_like])
vz = vz * -(m1[best_like]) / np.median(m0[best_like])
# plot rv trend
plt.plot(Time(epochs_seppa[i, :], format='mjd').decimalyear,
vz,
color=sep_pa_color)
# 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
def test_compute_model():
"""
Tests that the perturbation of a second planet using compute model gives roughly the amplitude we expect.
"""
# generate planet b orbital parameters
b_params = [1, 0, 0, 0, 0, 0]
tau_ref_epoch = 0
mass_b = 0.001 # Msun
m0 = 1 # Msun
plx = 1 # mas
# generate planet c orbital parameters
# at 0.3 au, and starts on the opposite side of the star relative to b
c_params = [0.3, 0, 0, np.pi, 0, 0]
mass_c = 0.002 # Msun
mtot = m0 + mass_b + mass_c
period_c = np.sqrt(c_params[0]**3 / mtot)
period_b = np.sqrt(b_params[0]**3 / mtot)
epochs = np.linspace(0, period_c * 365.25,
100) + tau_ref_epoch # the full period of c, MJD
ra_model, dec_model, vz_model = kepler.calc_orbit(
epochs,
b_params[0],
b_params[1],
b_params[2],
b_params[3],
b_params[4],
b_params[5],
plx,
mtot,
tau_ref_epoch=tau_ref_epoch)
# generate some fake measurements just to feed into system.py to test bookkeeping
# just make a 1 planet fit for now
t = table.Table([
epochs,
np.ones(epochs.shape, dtype=int), ra_model,
np.zeros(ra_model.shape), dec_model,
np.zeros(dec_model.shape)
],
names=[
"epoch", "object", "raoff", "raoff_err", "decoff",
"decoff_err"
])
filename = os.path.join(orbitize.DATADIR,
"multiplanet_fake_1planettest.csv")
t.write(filename, overwrite=True)
# create the orbitize system and generate model predictions using the standard 1 body model for b, and the 2 body model for b and c
astrom_dat = read_input.read_file(filename)
sys_1body = system.System(1,
astrom_dat,
m0,
plx,
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
sys_2body = system.System(2,
astrom_dat,
m0,
plx,
tau_ref_epoch=tau_ref_epoch,
fit_secondary_mass=True)
# model predictions for the 1 body case
# we had put one measurements of planet b in the data table, so compute_model only does planet b measurements
params = np.append(b_params, [plx, mass_b, m0])
radec_1body, _ = sys_1body.compute_model(params)
ra_1body = radec_1body[:, 0]
dec_1body = radec_1body[:, 1]
# model predictions for the 2 body case
# still only generates predictions of b's location, but including the perturbation for c
params = np.append(b_params, np.append(c_params,
[plx, mass_b, mass_c, m0]))
radec_2body, _ = sys_2body.compute_model(params)
ra_2body = radec_2body[:, 0]
dec_2body = radec_2body[:, 1]
ra_diff = ra_2body - ra_1body
dec_diff = dec_2body - dec_1body
total_diff = np.sqrt(ra_diff**2 + dec_diff**2)
# the expected influence of c is mass_c/m0 * sma_c * plx in amplitude
# just test the first value, because of the face on orbit, we should see it immediately.
assert total_diff[0] == pytest.approx(mass_c / m0 * c_params[0] * plx,
abs=0.01 * mass_c / m0 *
b_params[0] * plx)
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):
"""
Plots one orbital period for a select number of fitted orbits
for a given object, with line segments colored according to time
Args:
parallax (float): parallax (in mas), however, if plx_err was passed
to system, then this is ignored and the posterior samples for
plx will be used instead (default: None)
total_mass (float): total mass of system in solar masses, however,
if mass_err was passed to system, then this is ignored and the
posterior samples for mtot will be used instead (default: None)
object_mass (float): mass of the object, in solar masses (default: 0)
Note: this input has no effect at this time
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)
Return:
``matplotlib.pyplot.Figure``: the orbit plot if input is valid, None otherwise
(written): Henry Ngo, Sarah Blunt, 2018
"""
# Split the 2-D post array into series of 1-D arrays for each orbital parameter
num_objects, remainder = np.divmod(self.post.shape[1],6)
if object_to_plot > num_objects:
return None
first_index = 0 + 6*(object_to_plot-1)
sma = self.post[:,first_index+0]
ecc = self.post[:,first_index+1]
inc = self.post[:,first_index+2]
aop = self.post[:,first_index+3]
pan = self.post[:,first_index+4]
tau = self.post[:,first_index+5]
# Then, get the other parameters
if remainder == 2: # have samples for parallax and mtot
plx = self.post[:,-2]
mtot = self.post[:,-1]
else: # otherwise make arrays out of user provided value
if total_mass is not None:
mtot = np.ones(len(sma))*total_mass
else:
raise Exception('results.Results.plot_orbits(): total mass must be provided if not part of samples')
if parallax is not None:
plx = np.ones(len(sma))*parallax
else:
raise Exception('results.Results.plot_orbits(): parallax must be provided if not part of samples')
mplanet = np.ones(len(sma))*object_mass
# Select random indices for plotted orbit
if num_orbits_to_plot > len(sma):
num_orbits_to_plot = len(sma)
choose = np.random.randint(0, high=len(sma), size=num_orbits_to_plot)
raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
deoff = 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):
orb_ind = choose[i]
# 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[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],
tau[orb_ind], plx[orb_ind], mtot[orb_ind], mass=mplanet[orb_ind]
)
raoff[i,:] = raoff0
deoff[i,:] = deoff0
# Create a linearly increasing colormap for our range of epochs
norm = mpl.colors.Normalize(vmin=np.min(epochs), vmax=np.max(epochs[-1,:]))
norm_yr = mpl.colors.Normalize(vmin=np.min(epochs/365.25), vmax=np.max(epochs[-1,:]/365.25))
# Create figure for orbit plots
fig, ax = plt.subplots()
# 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
)
lc.set_array(epochs[i,:])
ax.add_collection(lc)
# 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)
# 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')
return fig
