def test_OFTI_multiplanet():
# initialize sampler
input_file = os.path.join(orbitize.DATADIR, "test_val_multi.csv")
myDriver = orbitize.driver.Driver(input_file,
'OFTI',
2,
1.52,
24.76,
mass_err=0.15,
plx_err=0.64)
s = myDriver.sampler
# change eccentricity prior for b
myDriver.system.sys_priors[1] = priors.UniformPrior(0.0, 0.1)
# change eccentricity prior for c
myDriver.system.sys_priors[7] = priors.UniformPrior(0.0, 0.1)
orbits = s.run_sampler(500)
idx = s.system.param_idx
sma1 = np.median(orbits[:, idx['sma1']])
sma2 = np.median(orbits[:, idx['sma2']])
sma1_exp = 66
sma2_exp = 40
print(sma1, sma2)
assert sma1 == pytest.approx(sma1_exp, abs=0.3 * sma1_exp)
assert sma2 == pytest.approx(sma2_exp, abs=0.3 * sma2_exp)
assert np.all(orbits[:, idx['ecc1']] < 0.1)
assert np.all(orbits[:, idx['ecc2']] < 0.1)
def main():
#parameters for the system
num_planets = 1
data_table = read_input.read_file('../Data/data.csv')
m0 = 1.01
mass_err = 0.05
plx = 78.33591471044681
plx_err = 0.1
#initialise a system object
sys = system.System(num_planets,
data_table,
m0,
plx,
mass_err=mass_err,
plx_err=plx_err,
fit_secondary_mass=True)
sys.sys_priors[lab['plx1']] = priors.UniformPrior(60, 110)
sys.sys_priors[lab['sma1']] = priors.UniformPrior(0.5, 1.50)
#MCMC parameters
n_temps = 10
n_walkers = 1000
n_threads = 10
total_orbits_MCMC = 75000000
burn_steps = 15000
thin = 10
#set up sampler object and run it
mcmc_sampler = sampler.MCMC(sys, n_temps, n_walkers, n_threads)
orbits = mcmc_sampler.run_sampler(total_orbits_MCMC,
burn_steps=burn_steps,
thin=thin)
#save results
myResults = mcmc_sampler.results
try:
### CHANGE THIS TO SAVE TO YOUR DESIRED DIRECTORY ##
filename = 'floatplx.hdf5'
# hdf5_filename=os.path.join(save_path,filename)
myResults.save_results(
filename) # saves results object as an hdf5 file
except:
print("Something went wrong while saving the results")
finally:
corner_figure = myResults.plot_corner()
corner_name = 'floatplx_corner.png'
corner_figure.savefig(corner_name)
orbit_figure = myResults.plot_orbits(rv_time_series=True)
orbit_name = 'floatplx_orbit.png'
orbit_figure.savefig(orbit_name)
return None
def test_OFTI_pan_priors():
# initialize sampler
myDriver = orbitize.driver.Driver(input_file,
'OFTI',
1,
1.22,
56.95,
mass_err=0.08,
plx_err=0.26)
s = myDriver.sampler
# change PAN prior
new_min = 0.05
new_max = np.pi - 0.05
myDriver.system.sys_priors[4] = priors.UniformPrior(new_min, new_max)
# run sampler
orbits = s.run_sampler(100)
# check that bounds were applied correctly
assert np.max(orbits[:, 4]) < new_max
assert np.min(orbits[:, 4]) > new_min
# change PAN prior again
mu = np.pi / 2
sigma = 0.05
myDriver.system.sys_priors[4] = priors.GaussianPrior(mu, sigma=sigma)
# run sampler again
orbits = s.run_sampler(250)
# check that bounds were applied correctly
assert mu == pytest.approx(np.mean(orbits[:, 4]), abs=0.01)
assert sigma == pytest.approx(np.std(orbits[:, 4]), abs=0.01)
def __init__(self,
num_secondary_bodies,
data_table,
system_mass,
plx,
mass_err=0,
plx_err=0,
restrict_angle_ranges=None,
results=None):
self.num_secondary_bodies = num_secondary_bodies
self.sys_priors = []
self.labels = []
self.results = []
#
# Group the data in some useful ways
#
self.data_table = data_table
# 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 (len(radec_indices) + len(seppa_indices) == len(
self.data_table)) and (restrict_angle_ranges is None):
restrict_angle_ranges = True
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.JeffreysPrior(0.1, 100.))
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., angle_upperlim))
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('epp{}'.format(body + 1))
#
# Set priors on total mass and parallax
#
self.labels.append('plx')
self.labels.append('mtot')
if plx_err > 0:
self.sys_priors.append(priors.GaussianPrior(plx, plx_err))
else:
self.sys_priors.append(plx)
if mass_err > 0:
self.sys_priors.append(priors.GaussianPrior(system_mass, mass_err))
else:
self.sys_priors.append(system_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 = []
# 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))))
else:
fit_secondary_mass = False
betaPic_Hip = None
betaPic_gaia = None
betaPic_system = system.System(num_secondary_bodies,
data_table,
1.75,
plx,
hipparcos_IAD=betaPic_Hip,
gaia=betaPic_gaia,
fit_secondary_mass=fit_secondary_mass,
mass_err=0.01,
plx_err=0.01)
m0_or_mtot_prior = priors.UniformPrior(1.5, 2.0)
# set uniform parallax prior
plx_index = betaPic_system.param_idx['plx']
betaPic_system.sys_priors[plx_index] = priors.UniformPrior(
plx - 1.0, plx + 1.0)
# set prior on Omega, since we know that know direction of orbital motion from RV
pan_index = betaPic_system.param_idx['pan1']
betaPic_system.sys_priors[pan_index] = priors.UniformPrior(0, np.pi)
# set uniform prior on m1 as Nielsen+ 2020 do
m1_index = betaPic_system.param_idx['m1']
betaPic_system.sys_priors[m1_index] = priors.UniformPrior(0, 0.1)
if fit_IAD:
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))))