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 __init__(self,
num_secondary_bodies,
data_table,
stellar_mass,
plx,
mass_err=0,
plx_err=0,
restrict_angle_ranges=None,
tau_ref_epoch=58849,
fit_secondary_mass=False,
results=None):
self.num_secondary_bodies = num_secondary_bodies
self.sys_priors = []
self.labels = []
self.results = []
self.fit_secondary_mass = fit_secondary_mass
self.tau_ref_epoch = tau_ref_epoch
#
# Group the data in some useful ways
#
self.data_table = data_table
# Creates a copy of the input in case data_table needs to be modified
self.input_table = self.data_table.copy()
# List of arrays of indices corresponding to each body
self.body_indices = []
# List of arrays of indices corresponding to epochs in RA/Dec for each body
self.radec = []
# List of arrays of indices corresponding to epochs in SEP/PA for each body
self.seppa = []
# List of index arrays corresponding to each rv for each body
self.rv = []
radec_indices = np.where(self.data_table['quant_type'] == 'radec')
seppa_indices = np.where(self.data_table['quant_type'] == 'seppa')
rv_indices = np.where(self.data_table['quant_type'] == 'rv')
# save indicies for all of the ra/dec, sep/pa measurements for convenience
self.all_radec = radec_indices
self.all_seppa = seppa_indices
for body_num in np.arange(self.num_secondary_bodies + 1):
self.body_indices.append(
np.where(self.data_table['object'] == body_num))
self.radec.append(
np.intersect1d(self.body_indices[body_num], radec_indices))
self.seppa.append(
np.intersect1d(self.body_indices[body_num], seppa_indices))
self.rv.append(
np.intersect1d(self.body_indices[body_num], rv_indices))
# we should track the influence of the planet(s) on each other/the star if we are not fitting massless planets and
# we are not fitting relative astrometry of just a single body
self.track_planet_perturbs = self.fit_secondary_mass and \
((len(self.radec[1]) + len(self.seppa[1]) + len(self.rv[1]) < len(data_table)) or \
(self.num_secondary_bodies > 1))
if restrict_angle_ranges:
angle_upperlim = np.pi
else:
angle_upperlim = 2. * np.pi
#
# Set priors for each orbital element
#
for body in np.arange(num_secondary_bodies):
# Add semimajor axis prior
self.sys_priors.append(priors.LogUniformPrior(0.001, 1e7))
self.labels.append('sma{}'.format(body + 1))
# Add eccentricity prior
self.sys_priors.append(priors.UniformPrior(0., 1.))
self.labels.append('ecc{}'.format(body + 1))
# Add inclination angle prior
self.sys_priors.append(priors.SinPrior())
self.labels.append('inc{}'.format(body + 1))
# Add argument of periastron prior
self.sys_priors.append(priors.UniformPrior(0., 2. * np.pi))
self.labels.append('aop{}'.format(body + 1))
# Add position angle of nodes prior
self.sys_priors.append(priors.UniformPrior(0., angle_upperlim))
self.labels.append('pan{}'.format(body + 1))
# Add epoch of periastron prior.
self.sys_priors.append(priors.UniformPrior(0., 1.))
self.labels.append('tau{}'.format(body + 1))
#
# Set priors on total mass and parallax
#
self.labels.append('plx')
if plx_err > 0:
self.sys_priors.append(priors.GaussianPrior(plx, plx_err))
else:
self.sys_priors.append(plx)
# checking for rv data to include appropriate rv priors:
if len(self.rv[0]) > 0 and self.fit_secondary_mass:
self.sys_priors.append(priors.UniformPrior(
-5, 5)) # gamma prior in km/s
self.labels.append('gamma')
self.sys_priors.append(priors.LogUniformPrior(
1e-4, 0.05)) # jitter prior in km/s
self.labels.append('sigma')
if self.fit_secondary_mass:
for body in np.arange(num_secondary_bodies) + 1:
self.sys_priors.append(priors.LogUniformPrior(
1e-6, 2)) # in Solar masses for now
self.labels.append('m{}'.format(body))
self.labels.append('m0')
else:
self.labels.append('mtot')
# still need to append m0/mtot, even though labels are appended above
if mass_err > 0:
self.sys_priors.append(priors.GaussianPrior(
stellar_mass, mass_err))
else:
self.sys_priors.append(stellar_mass)
# add labels dictionary for parameter indexing
self.param_idx = dict(zip(self.labels, np.arange(len(self.labels))))
def __init__(self,
num_secondary_bodies,
data_table,
stellar_mass,
plx,
mass_err=0,
plx_err=0,
restrict_angle_ranges=None,
tau_ref_epoch=58849,
fit_secondary_mass=False,
results=None):
self.num_secondary_bodies = num_secondary_bodies
self.sys_priors = []
self.labels = []
self.results = []
self.fit_secondary_mass = fit_secondary_mass
self.tau_ref_epoch = tau_ref_epoch
self.restrict_angle_ranges = restrict_angle_ranges
#
# Group the data in some useful ways
#
self.data_table = data_table
# Creates a copy of the input in case data_table needs to be modified
self.input_table = self.data_table.copy()
# Rob: check if instrument column is other than default. If other than default, then separate data table into n number of instruments.
# gather list of instrument: list_instr = self.data_table['instruments'][name of instrument]
# List of arrays of indices corresponding to each body
# instruments = np.unique(self.data_table['instruments']) gives a list of unique names
self.body_indices = []
# List of arrays of indices corresponding to epochs in RA/Dec for each body
self.radec = []
# List of arrays of indices corresponding to epochs in SEP/PA for each body
self.seppa = []
# List of index arrays corresponding to each rv for each body
self.rv = []
# instr1_tbl = np.where(self.data_table['instruments'] == list_instr[0])
# loop through the indices per input_table:
# rv_indices = np.where(instr1_tbl['quant_type'] == 'rv')
# ... return the parameter labels for each table
# ...
self.fit_astrometry = True
radec_indices = np.where(self.data_table['quant_type'] == 'radec')
seppa_indices = np.where(self.data_table['quant_type'] == 'seppa')
if len(radec_indices[0]) == 0 and len(seppa_indices[0]) == 0:
self.fit_astrometry = False
rv_indices = np.where(self.data_table['quant_type'] == 'rv')
# defining all indices to loop through the unique rv instruments to get different offsets and jitters
instrument_list = np.unique(self.data_table['instrument'])
inst_indices_all = []
for inst in instrument_list:
inst_indices = np.where(self.data_table['instrument'] == inst)
inst_indices_all.append(inst_indices)
# defining indices for unique instruments in the data table
self.rv_instruments = np.unique(
self.data_table['instrument'][rv_indices])
self.rv_inst_indices = []
for inst in self.rv_instruments:
inst_indices = np.where(self.data_table['instrument'] == inst)
self.rv_inst_indices.append(inst_indices)
# astrometry instruments same for radec and seppa:
self.astr_instruments = np.unique(
self.data_table['instrument'][np.where(
self.data_table['quant_type'] != 'rv')])
# save indicies for all of the ra/dec, sep/pa measurements for convenience
self.all_radec = radec_indices
self.all_seppa = seppa_indices
for body_num in np.arange(self.num_secondary_bodies + 1):
self.body_indices.append(
np.where(self.data_table['object'] == body_num))
self.radec.append(
np.intersect1d(self.body_indices[body_num], radec_indices))
self.seppa.append(
np.intersect1d(self.body_indices[body_num], seppa_indices))
self.rv.append(
np.intersect1d(self.body_indices[body_num], rv_indices))
# we should track the influence of the planet(s) on each other/the star if we are not fitting massless planets and
# we are not fitting relative astrometry of just a single body
self.track_planet_perturbs = self.fit_secondary_mass and \
((len(self.radec[1]) + len(self.seppa[1]) + len(self.rv[1]) < len(data_table)) or \
(self.num_secondary_bodies > 1))
if restrict_angle_ranges:
angle_upperlim = np.pi
else:
angle_upperlim = 2. * np.pi
#
# Set priors for each orbital element
#
for body in np.arange(num_secondary_bodies):
# Add semimajor axis prior
self.sys_priors.append(priors.LogUniformPrior(0.001, 1e7))
self.labels.append('sma{}'.format(body + 1))
# Add eccentricity prior
self.sys_priors.append(priors.UniformPrior(0., 1.))
self.labels.append('ecc{}'.format(body + 1))
# Add inclination angle prior
self.sys_priors.append(priors.SinPrior())
self.labels.append('inc{}'.format(body + 1))
# Add argument of periastron prior
self.sys_priors.append(priors.UniformPrior(0., 2. * np.pi))
self.labels.append('aop{}'.format(body + 1))
# Add position angle of nodes prior
self.sys_priors.append(priors.UniformPrior(0., angle_upperlim))
self.labels.append('pan{}'.format(body + 1))
# Add epoch of periastron prior.
self.sys_priors.append(priors.UniformPrior(0., 1.))
self.labels.append('tau{}'.format(body + 1))
#
# Set priors on total mass and parallax
#
self.labels.append('plx')
if plx_err > 0:
self.sys_priors.append(priors.GaussianPrior(plx, plx_err))
else:
self.sys_priors.append(plx)
# checking for rv data to include appropriate rv priors:
if len(self.rv[0]) > 0 and self.fit_secondary_mass:
# Rob and Lea:
# for instrument in rv_instruments:
# add gamma and sigma for each and label each unique gamma and sigma per instrument name (gamma+instrument1, ...)
for instrument in self.rv_instruments:
self.sys_priors.append(priors.UniformPrior(
-5, 5)) # gamma prior in km/s
self.labels.append('gamma_{}'.format(instrument))
self.sys_priors.append(priors.LogUniformPrior(
1e-4, 0.05)) # jitter prior in km/s
self.labels.append('sigma_{}'.format(instrument))
if self.fit_secondary_mass:
for body in np.arange(num_secondary_bodies) + 1:
self.sys_priors.append(priors.LogUniformPrior(
1e-6, 2)) # in Solar masses for now
self.labels.append('m{}'.format(body))
self.labels.append('m0')
else:
self.labels.append('mtot')
# still need to append m0/mtot, even though labels are appended above
if mass_err > 0:
self.sys_priors.append(priors.GaussianPrior(
stellar_mass, mass_err))
else:
self.sys_priors.append(stellar_mass)
# add labels dictionary for parameter indexing
self.param_idx = dict(zip(self.labels, np.arange(len(self.labels))))
def __init__(self,
num_secondary_bodies,
data_table,
stellar_mass,
plx,
mass_err=0,
plx_err=0,
restrict_angle_ranges=None,
tau_ref_epoch=58849,
fit_secondary_mass=False,
results=None):
self.num_secondary_bodies = num_secondary_bodies
self.sys_priors = []
self.labels = []
self.results = []
self.fit_secondary_mass = fit_secondary_mass
self.tau_ref_epoch = tau_ref_epoch
#
# Group the data in some useful ways
#
self.data_table = data_table
# Creates a copy of the input in case data_table needs to be modified
self.input_table = self.data_table.copy()
# List of arrays of indices corresponding to each body
self.body_indices = []
# List of arrays of indices corresponding to epochs in RA/Dec for each body
self.radec = []
# List of arrays of indices corresponding to epochs in SEP/PA for each body
self.seppa = []
radec_indices = np.where(self.data_table['quant_type'] == 'radec')
seppa_indices = np.where(self.data_table['quant_type'] == 'seppa')
for body_num in np.arange(self.num_secondary_bodies + 1):
self.body_indices.append(
np.where(self.data_table['object'] == body_num))
self.radec.append(
np.intersect1d(self.body_indices[body_num], radec_indices))
self.seppa.append(
np.intersect1d(self.body_indices[body_num], seppa_indices))
if restrict_angle_ranges:
angle_upperlim = np.pi
else:
angle_upperlim = 2. * np.pi
#
# Set priors for each orbital element
#
for body in np.arange(num_secondary_bodies):
# Add semimajor axis prior
self.sys_priors.append(priors.LogUniformPrior(0.001, 1e7))
self.labels.append('sma{}'.format(body + 1))
# Add eccentricity prior
self.sys_priors.append(priors.UniformPrior(0., 1.))
self.labels.append('ecc{}'.format(body + 1))
# Add inclination angle prior
self.sys_priors.append(priors.SinPrior())
self.labels.append('inc{}'.format(body + 1))
# Add argument of periastron prior
self.sys_priors.append(priors.UniformPrior(0., 2. * np.pi))
self.labels.append('aop{}'.format(body + 1))
# Add position angle of nodes prior
self.sys_priors.append(priors.UniformPrior(0., angle_upperlim))
self.labels.append('pan{}'.format(body + 1))
# Add epoch of periastron prior.
self.sys_priors.append(priors.UniformPrior(0., 1.))
self.labels.append('tau{}'.format(body + 1))
#
# Set priors on total mass and parallax
#
self.labels.append('plx')
if plx_err > 0:
self.sys_priors.append(priors.GaussianPrior(plx, plx_err))
else:
self.sys_priors.append(plx)
if self.fit_secondary_mass:
for body in np.arange(num_secondary_bodies):
self.sys_priors.append(priors.LogUniformPrior(
1e-6, 1)) # in Solar masses for now
self.labels.append('m{}'.format(body))
self.labels.append('m0')
else:
self.labels.append('mtot')
# still need to append m0/mtot, even though labels are appended above
if mass_err > 0:
self.sys_priors.append(priors.GaussianPrior(
stellar_mass, mass_err))
else:
self.sys_priors.append(stellar_mass)
# add labels dictionary for parameter indexing
self.param_idx = dict(zip(self.labels, np.arange(len(self.labels))))