def test_custom_likelihood():
"""
Tests the inclusion of a custom likelihood function in the code
"""
# use the test_csv dir
testdir = orbitize.DATADIR
input_file = os.path.join(testdir, 'GJ504.csv')
data_table = read_input.read_file(input_file)
# Manually set 'object' column of data table
data_table['object'] = 1
# construct the system
orbit = system.System(1, data_table, 1, 0.01)
# construct custom likelihood function
def my_likelihood(params):
return -5
# construct sampler
n_walkers = 100
mcmc1 = sampler.MCMC(orbit, 0, n_walkers, num_threads=1)
mcmc2 = sampler.MCMC(
orbit, 0, n_walkers, num_threads=1, custom_lnlike=my_likelihood
)
param = np.array([2, 0.5, 0.5, 0.5, 0.5, 0.5, 1, 0.01])
logl1 = mcmc1._logl(param)
logl2 = mcmc2._logl(param)
assert logl1 == logl2 + 5
def test_pt_mcmc_runs(num_threads=1):
"""
Tests the PTMCMC sampler by making sure it even runs
"""
# use the test_csv dir
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
data_table = read_input.read_formatted_file(input_file)
# Manually set 'object' column of data table
data_table['object'] = 1
# construct the system
orbit = system.System(1, data_table, 1, 0.01)
# construct sampler
n_temps = 2
n_walkers = 100
mcmc = sampler.MCMC(orbit, n_temps, n_walkers, num_threads=num_threads)
# run it a little (tests 0 burn-in steps)
mcmc.run_sampler(100)
# run it a little more
mcmc.run_sampler(1000, 1)
# run it a little more (tests adding to results object)
mcmc.run_sampler(1000, 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 run_fit(fname):
'''
Runs the orbit fit! Saves the resultant posterior, orbit plot and corner plot
args:
fname (str): Path to the data file.
'''
#parameters for the system
num_planets=1
data_table = read_input.read_file(fname)
m0 = 1.0
mass_err = 0.01
plx=50
plx_err=0.01
#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
)
#MCMC parameters
n_temps=5
n_walkers=1000
n_threads=5
total_orbits_MCMC=5000 # n_walkers x num_steps_per_walker
burn_steps=1
thin=1
#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)
myResults=mcmc_sampler.results
try:
save_path = '.'
filename = 'post.hdf5'
hdf5_filename=os.path.join(save_path,filename)
myResults.save_results(hdf5_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='corner.png'
corner_figure.savefig(corner_name)
orbit_figure=myResults.plot_orbits(rv_time_series=True)
orbit_name='joint_orbit.png'
orbit_figure.savefig(orbit_name)
print("Done!")
def test_save_and_load_gaia_and_hipparcos():
"""
Test that a Results object for a Gaia+Hipparcos fit
is saved and loaded properly.
"""
hip_num = '027321'
gaia_num = 4792774797545105664
num_secondary_bodies = 1
path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num)
myHip = hipparcos.HipparcosLogProb(path_to_iad_file, hip_num,
num_secondary_bodies)
myGaia = gaia.GaiaLogProb(gaia_num, myHip)
input_file = os.path.join(DATADIR, 'betaPic.csv')
data_table_with_rvs = read_input.read_file(input_file)
mySys = system.System(1,
data_table_with_rvs,
1.22,
56.95,
mass_err=0.08,
plx_err=0.26,
hipparcos_IAD=myHip,
gaia=myGaia,
fit_secondary_mass=True)
mySamp = sampler.MCMC(mySys, num_temps=1, num_walkers=50)
mySamp.run_sampler(1, burn_steps=0)
save_name = 'test_results.h5'
mySamp.results.save_results(save_name)
loadedResults = results.Results()
loadedResults.load_results(save_name)
assert np.all(loadedResults.system.hipparcos_IAD.epochs ==
mySys.hipparcos_IAD.epochs)
assert np.all(loadedResults.system.tau_ref_epoch == mySys.tau_ref_epoch)
assert np.all(loadedResults.system.gaia.ra == mySys.gaia.ra)
os.system('rm {}'.format(save_name))
# set uniform primary mass prior
m0_index = HD142527_system.param_idx['m0']
HD142527_system.sys_priors[m0_index] = priors.GaussianPrior(2.05, 0.3)
# set uniform primary mass prior
m1_index = HD142527_system.param_idx['m1']
HD142527_system.sys_priors[m1_index] = priors.GaussianPrior(0.25, 0.2)
# MCMC parameters
num_temps = 20
num_walkers = 50
num_threads = int(mp.cpu_count() * 2 / 3)
# init driver
HD142527_sampler = sampler.MCMC(HD142527_system,
num_threads=num_threads,
num_temps=num_temps,
num_walkers=num_walkers)
print('Driver Initialized! Running MCMC...')
HD142527_sampler.run_sampler(total_orbits,
burn_steps=burn_steps,
thin=thin)
myResults = HD142527_sampler.results
# save posterior to disk
savefile = 'paperdraft_posterior.hdf5'
myResults.save_results(savefile)
print('Saved Posterior!')
starttime = Time(datetime.strptime('1990 January 1',
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_examine_chop_chains(num_temps=0, num_threads=1):
"""
Tests the MCMC sampler's examine_chains and chop_chains methods
Args:
num_temps: Number of temperatures to use
Uses Parallel Tempering MCMC (ptemcee) if > 1,
otherwises, uses Affine-Invariant Ensemble Sampler (emcee)
num_threads: number of threads to run
"""
# use the test_csv dir
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
data_table = read_input.read_formatted_file(input_file)
# Manually set 'object' column of data table
data_table['object'] = 1
# construct the system
orbit = system.System(1, data_table, 1, 0.01)
# construct Driver
n_walkers = 20
mcmc = sampler.MCMC(orbit, num_temps, n_walkers, num_threads=num_threads)
# run it a little
n_samples1 = 2000 # 100 steps for each of 20 walkers
n_samples2 = 2000 # 100 steps for each of 20 walkers
n_samples = n_samples1 + n_samples2
mcmc.run_sampler(n_samples1)
# run it a little more (tries examine_chains within run_sampler)
mcmc.run_sampler(n_samples2, examine_chains=True)
# (4000 orbit samples = 20 walkers x 200 steps)
# Try all variants of examine_chains
mcmc.examine_chains()
plt.close('all') # Close figures generated
fig_list = mcmc.examine_chains(param_list=['sma1', 'ecc1', 'inc1'])
# Should only get 3 figures
assert len(fig_list) == 3
plt.close('all') # Close figures generated
mcmc.examine_chains(walker_list=[10, 12])
plt.close('all') # Close figures generated
mcmc.examine_chains(n_walkers=5)
plt.close('all') # Close figures generated
mcmc.examine_chains(step_range=[50, 100])
plt.close('all') # Close figures generated
# Now try chopping the chains
# Chop off first 50 steps
chop1 = 50
mcmc.chop_chains(chop1)
# Calculate expected number of orbits now
expected_total_orbits = n_samples - chop1 * n_walkers
# Check lengths of arrays in results object
assert len(mcmc.results.lnlike) == expected_total_orbits
assert mcmc.results.post.shape[0] == expected_total_orbits
# With 150 steps left, now try to trim 25 steps off each end
chop2 = 25
trim2 = 25
mcmc.chop_chains(chop2, trim=trim2)
# Calculated expected number of orbits now
samples_removed = (chop1 + chop2 + trim2) * n_walkers
expected_total_orbits = n_samples - samples_removed
# Check lengths of arrays in results object
assert len(mcmc.results.lnlike) == expected_total_orbits
assert mcmc.results.post.shape[0] == expected_total_orbits
sysrv=sysrv,
sysrv_err=sysrv_err,
mcmc_kwargs={
'num_temps': num_temps,
'num_walkers': num_walkers,
'num_threads': num_threads
},
system_kwargs={"restrict_angle_ranges": restrict_angle_ranges},
)
if len(planet) == 1:
my_driver.system.sys_priors[0] = priors.JeffreysPrior(1, 1e2)
my_driver.sampler = sampler.MCMC(my_driver.system,
num_temps=num_temps,
num_walkers=num_walkers,
num_threads=num_threads,
like='chi2_lnlike',
custom_lnlike=None)
if 1 and len(planet) == 1:
_filename = os.path.join(
astrometry_DATADIR, "figures", "HR_8799_" + planet,
'chain_' + rv_str + '_' + planet +
'_it{0}_{1}_{2}_{3}_{4}_{5}_coplanar.fits'.format(
itnum - 1, num_temps, num_walkers, 100000, thin, uservs))
print(_filename)
with pyfits.open(_filename) as hdulist: #total_orbits//num_walkers
print(hdulist[0].data.shape)
my_driver.sampler.curr_pos = hdulist[0].data[:, :, -1, :]
my_driver.sampler.save_intermediate = os.path.join(
tau_ref_epoch = 50000
# Read in data
data_table = read_input.read_file(datafile)
# Initialize System object which stores data & sets priors
my_system = system.System(num_secondary_bodies,
data_table,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
tau_ref_epoch=tau_ref_epoch)
my_sampler = sampler.MCMC(my_system, n_temps, n_walkers, n_threads)
# Run the sampler to compute some orbits, yeah!
# Results stored in bP_sampler.chain and bP_sampler.lnlikes
my_sampler.run_sampler(total_orbits, burn_steps=burn_steps, thin=10)
my_sampler.results.save_results("hr8799e_gravity_chains.hdf5")
#import orbitize.results as results
#my_sampler.results = results.Results()
#my_sampler.results.load_results("hr8799e_gravity_chains.hdf5")
# make corner plot
fig = my_sampler.results.plot_corner()
plt.savefig('corner_hr8799e_gravity.png', dpi=250)
def test_mcmc_runs(num_temps=0, num_threads=1):
"""
Tests the MCMC sampler by making sure it even runs
Args:
num_temps: Number of temperatures to use
Uses Parallel Tempering MCMC (ptemcee) if > 1,
otherwises, uses Affine-Invariant Ensemble Sampler (emcee)
num_threads: number of threads to run
"""
# use the test_csv dir
input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv')
data_table = read_input.read_file(input_file)
# Manually set 'object' column of data table
data_table['object'] = 1
# construct Driver
n_walkers = 100
myDriver = Driver(input_file,
'MCMC',
1,
1,
0.01,
mcmc_kwargs={
'num_temps': num_temps,
'num_threads': num_threads,
'num_walkers': n_walkers
})
# run it a little (tests 0 burn-in steps)
myDriver.sampler.run_sampler(100)
assert myDriver.sampler.results.post.shape[0] == 100
# run it a little more
myDriver.sampler.run_sampler(1000, burn_steps=1)
assert myDriver.sampler.results.post.shape[0] == 1100
# run it a little more (tests adding to results object, and periodic saving)
output_filename = os.path.join(orbitize.DATADIR, 'test_mcmc.hdf5')
myDriver.sampler.run_sampler(400,
burn_steps=1,
output_filename=output_filename,
periodic_save_freq=2)
# test results object exists and has 2100*100 steps
assert os.path.exists(output_filename)
saved_results = results.Results()
saved_results.load_results(output_filename)
assert saved_results.post.shape[0] == 1500
assert saved_results.curr_pos is not None # current positions should be saved
assert np.all(saved_results.curr_pos == myDriver.sampler.curr_pos)
# also check it is consistent with the internal results object in myDriver
assert myDriver.sampler.results.post.shape[0] == 1500
# run it a little more testing that everything gets saved even if prediodic_save_freq is not a multiple of the number of steps
output_filename_2 = os.path.join(orbitize.DATADIR, 'test_mcmc_v1.hdf5')
myDriver.sampler.run_sampler(500,
burn_steps=1,
output_filename=output_filename_2,
periodic_save_freq=3)
assert myDriver.sampler.results.post.shape[0] == 2000
# test that lnlikes being saved are correct
returned_lnlike_test = myDriver.sampler.results.lnlike[0]
computed_lnlike_test = myDriver.sampler._logl(
myDriver.sampler.results.post[0])
assert returned_lnlike_test == pytest.approx(computed_lnlike_test,
abs=0.01)
# test resuming and restarting from a prevous save
new_sampler = sampler.MCMC(myDriver.system,
num_temps=num_temps,
num_walkers=n_walkers,
num_threads=num_threads,
prev_result_filename=output_filename)
assert new_sampler.results.post.shape[0] == 1500
new_sampler.run_sampler(500, burn_steps=1)
assert new_sampler.results.post.shape[0] == 2000
assert new_sampler.results.post[0, 0] == myDriver.sampler.results.post[0,
0]
assert not betaPic_system.track_planet_perturbs
# set uniform mtot prior
mtot_index = betaPic_system.param_idx['mtot']
betaPic_system.sys_priors[mtot_index] = m0_or_mtot_prior
# run MCMC
num_threads = 100
num_temps = 20
num_walkers = 1000
num_steps = 1000000 # n_walkers x n_steps_per_walker
burn_steps = 1000
thin = 100
betaPic_sampler = sampler.MCMC(betaPic_system,
num_threads=num_threads,
num_temps=num_temps,
num_walkers=num_walkers)
betaPic_sampler.run_sampler(num_steps, burn_steps=burn_steps, thin=thin)
# save chains
betaPic_sampler.results.save_results('{}/betaPic_IAD{}.hdf5'.format(
savedir, fit_IAD))
# make corner plot
fig = betaPic_sampler.results.plot_corner()
plt.savefig('{}/corner_IAD{}.png'.format(savedir, fit_IAD), dpi=250)
# make orbit plot
fig = betaPic_sampler.results.plot_orbits()
plt.savefig('{}/orbit_IAD{}.png'.format(savedir, fit_IAD), dpi=250)
