def test_read_file():
"""
Test the read_file function using the test_val.csv file and test_val_radec.csv
"""
# Check that main test input is read in with correct values
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
_compare_table(read_file(input_file))
# Check that an input value with all valid entries and only ra/dec columns can be read
input_file_radec = os.path.join(orbitize.DATADIR, 'test_val_radec.csv')
read_file(input_file_radec)
def test_read_file():
"""
Test the read_file function using the test_val.csv file and test_val_radec.csv
"""
testdir = os.path.dirname(os.path.abspath(__file__))
# Check that main test input is read in with correct values
input_file = os.path.join(testdir, 'test_val.csv')
_compare_table(read_file(input_file))
# Check that an input value with all valid entries and only ra/dec columns can be read
input_file_radec = os.path.join(testdir, 'test_val_radec.csv')
read_file(input_file_radec)
def test_semi_amp_basis():
"""
For both MCMC and OFTI param formats, make the conversion to the standard
basis from semi-amplitude and back to verify the valdity of conversions. Do
this with a single companion and with two companions.
"""
# 1. Single Body (with RV)
filename = "{}/HD4747.csv".format(DATADIR)
data_table = read_input.read_file(filename)
my_system = system.System(
1, data_table, 0.84, 53.18, mass_err=0.04, plx_err=0.12,
fit_secondary_mass=True, fitting_basis='SemiAmp'
)
num_samples = 100
samples = np.empty([len(my_system.sys_priors), num_samples])
for i in range(len(my_system.sys_priors)):
if hasattr(my_system.sys_priors[i], "draw_samples"):
samples[i, :] = my_system.sys_priors[i].draw_samples(num_samples)
else:
samples[i, :] = my_system.sys_priors[i] * np.ones(num_samples)
sample_copy = samples.copy()
# MCMC Format
test = samples[:, 0].copy()
conversion = my_system.basis.to_standard_basis(test)
original = my_system.basis.to_semi_amp_basis(conversion)
assert np.allclose(original, sample_copy[:, 0])
# 2. Multi Body
filename = "{}/test_val_multi.csv".format(DATADIR)
data_table = read_input.read_file(filename)
my_system = system.System(
2, data_table, 1.52, 24.76, mass_err=0.15, plx_err=0.64,
fit_secondary_mass=True, fitting_basis='SemiAmp'
)
num_samples = 100
samples = np.empty([len(my_system.sys_priors), num_samples])
for i in range(len(my_system.sys_priors)):
if hasattr(my_system.sys_priors[i], "draw_samples"):
samples[i, :] = my_system.sys_priors[i].draw_samples(num_samples)
else:
samples[i, :] = my_system.sys_priors[i] * np.ones(num_samples)
sample_copy = samples.copy()
# MCMC Format
test = samples[:, 0].copy()
conversion = my_system.basis.to_standard_basis(test)
original = my_system.basis.to_semi_amp_basis(conversion)
assert np.allclose(original, sample_copy[:, 0])
def test_read_formatted_file():
"""
Tests the read_formatted_file function using the test_val.csv file and test_val_radec.csv
This test exists with the fail_if_not_removed decorator as a reminder to remove in v2.0
"""
# Check that main test input is read in with correct values
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
_compare_table(read_formatted_file(input_file))
# Check that an input value with all valid entries and only ra/dec columns can be read
input_file_radec = os.path.join(orbitize.DATADIR, 'test_val_radec.csv')
read_file(input_file_radec)
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_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 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_multi_planets():
"""
Test using some data for HR 8799 b and c
"""
num_secondary_bodies = 2
input_file = os.path.join(orbitize.DATADIR, 'test_val_multi.csv')
data_table = read_input.read_file(input_file)
system_mass = 1.47
plx = 24.30
mass_err = 0.11
plx_err = 0.7
# Initialize System object, use mjd=50000 to be consistent with Wang+2018
test_system = system.System(num_secondary_bodies,
data_table,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
tau_ref_epoch=50000)
params = np.array([
7.2774010e+01, 4.1116819e-02, 5.6322372e-01, 3.5251172e+00,
4.2904768e+00, 9.4234377e-02, 4.5418411e+01, 1.4317369e-03,
5.6322372e-01, 3.1016846e+00, 4.2904768e+00, 3.4033456e-01,
2.4589758e+01, 1.4849439e+00
])
result = test_system.compute_model(params)
print(result)
def test_add_and_clear_results():
num_secondary_bodies = 1
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
data_table = read_input.read_file(input_file)
system_mass = 1.0
plx = 10.0
mass_err = 0.1
plx_err = 1.0
# Initialize System object
test_system = system.System(num_secondary_bodies,
data_table,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err)
# Initialize dummy results.Results object
test_results = results.Results()
# Add one result object
test_system.add_results(test_results)
assert len(test_system.results) == 1
# Adds second result object
test_system.add_results(test_results)
assert len(test_system.results) == 2
# Clears result objects
test_system.clear_results()
assert len(test_system.results) == 0
# Add one more result object
test_system.add_results(test_results)
assert len(test_system.results) == 1
def test_systeminit():
"""
Test that initializing a ``System`` class produces a list of ``Prior``
objects of the correct length when:
- parallax and total mass are fixed
- parallax and total mass errors are given
- parallax is fixed, total mass error is given
- parallax error is given, total mass error is fixed
Test that the different types of data are parsed correctly
when initializing a ``System`` object.
"""
testdir = orbitize.DATADIR
input_file = os.path.join(testdir, 'test_val.csv')
data_table = read_input.read_file(input_file)
# Manually set 'object' column of data table
data_table['object'] = 1
data_table['object'][1] = 2
plx_mass_errs2lens = {
(0., 0.): 14,
(1., 1.): 14,
(0., 1.): 14,
(1., 0.): 14
}
for plx_e, mass_e in plx_mass_errs2lens.keys():
testSystem_priors = system.System(2,
data_table,
10.,
10.,
plx_err=plx_e,
mass_err=mass_e)
assert len(testSystem_priors.sys_priors) == \
plx_mass_errs2lens[(plx_e, mass_e)]
testSystem_parsing = system.System(2,
data_table,
10.,
10.,
plx_err=0.5,
mass_err=0.5)
assert len(data_table[testSystem_parsing.seppa[0]]) == 0
assert len(data_table[testSystem_parsing.seppa[1]]) == 1
assert len(data_table[testSystem_parsing.seppa[2]]) == 1
assert len(data_table[testSystem_parsing.radec[0]]) == 0
assert len(data_table[testSystem_parsing.radec[1]]) == 1
assert len(data_table[testSystem_parsing.radec[2]]) == 0
assert testSystem_parsing.labels == [
'sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'sma2', 'ecc2', 'inc2',
'aop2', 'pan2', 'tau2', 'plx', 'mtot'
]
def test_write_orbitize_input():
"""
Test the write_orbitize_input and the read_file functions
"""
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
test_table = read_file(input_file)
output_file = os.path.join(orbitize.DATADIR, 'temp_test_orbitize_input.csv')
# If temp output file already exists, delete it
if os.path.isfile(output_file):
os.remove(output_file)
try: # Catch these tests so that we remove temporary file
# Test that we were able to write the table
write_orbitize_input(test_table, output_file)
assert os.path.isfile(output_file)
# Test that we can read the table and check if it's correct
test_table_2 = read_file(output_file)
_compare_table(test_table_2)
finally:
# Remove temporary file
os.remove(output_file)
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_XYZ():
'''
Test the XYZ basis on data that does not throw exceptions using simulated data for
(1) where single companion is supplied and its mass is not being fitted, (2) a single
companion is supplied and its mass is being fitted, and (3) two companions are supplied
and their masses are being fitted.
'''
filename = '{}/xyz_test_data.csv'.format(DATADIR)
data = read_input.read_file(filename)
# (1) Single Companion (mtot)
single_data = data[np.where(data['object'] == 1)[0]]
copy = single_data.copy()
my_system = system.System(1,
single_data,
1.22,
56.89,
mass_err=0.05,
plx_err=0.12,
fitting_basis='XYZ')
expected_labels = [item + '1'
for item in basis_names['XYZ']] + ['plx', 'mtot']
assert expected_labels == my_system.labels
# (2) Single Companion (fit companion mass)
my_system = system.System(1,
copy,
1.22,
56.89,
mass_err=0.05,
plx_err=0.12,
fit_secondary_mass=True,
fitting_basis='XYZ')
expected_labels = [item + '1'
for item in basis_names['XYZ']] + ['plx', 'm1', 'm0']
assert expected_labels == my_system.labels
# (3) Two Companions (fit companion masses)
my_system = system.System(2,
data,
1.22,
56.89,
mass_err=0.05,
plx_err=0.12,
fit_secondary_mass=True,
fitting_basis='XYZ')
expected_labels = [item + '1' for item in basis_names['XYZ']] + [
item + '2' for item in basis_names['XYZ']
] + ['plx', 'm1', 'm2', 'm0']
assert expected_labels == my_system.labels
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_create_driver_from_table():
"""
Test creation of Driver object from Table as input
"""
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
input_table = read_file(input_file)
myDriver = driver.Driver(input_table, # astropy.table Table of input
'MCMC', # name of algorith for orbit-fitting
1, # number of secondary bodies in system
1.0, # total system mass [M_sun]
50.0, # total parallax of system [mas]
mass_err=0.1, # mass error [M_sun]
plx_err=0.1) # parallax error [mas]
_compare_table(myDriver.system.data_table)
def test_convert_data_table_radec2seppa():
num_secondary_bodies=1
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
data_table=read_input.read_file(input_file)
system_mass=1.0
plx=10.0
mass_err=0.1
plx_err=1.0
# Initialize System object
test_system = system.System(
num_secondary_bodies, data_table, system_mass,
plx, mass_err=mass_err, plx_err=plx_err
)
test_system.convert_data_table_radec2seppa()
def test_convert_data_table_radec2seppa():
num_secondary_bodies = 1
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
data_table = read_input.read_file(input_file)
system_mass = 1.0
plx = 10.0
mass_err = 0.1
plx_err = 1.0
# Initialize System object
test_system = system.System(num_secondary_bodies,
data_table,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err)
test_system.convert_data_table_radec2seppa()
def test_create_driver_from_table():
"""
Test creation of Driver object from Table as input
"""
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
input_table = read_file(input_file)
myDriver = driver.Driver(
input_table, # astropy.table Table of input
'MCMC', # name of algorithm for orbit-fitting
1, # number of secondary bodies in system
1.0, # total system mass [M_sun]
50.0, # total parallax of system [mas]
mass_err=0.1, # mass error [M_sun]
plx_err=0.1, # parallax error [mas]
system_kwargs={'fit_secondary_mass': True})
_compare_table(myDriver.system.data_table)
def test_compute_model():
"""
Test basic functionality of ``System.compute_model()``
"""
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
data_table = read_input.read_file(input_file)
data_table['object'] = 1
testSystem_parsing = system.System(1, data_table, 10., 10.)
params_arr = np.array([[1., 0.5], [0., 0.], [0., 0.], [0., 0.], [0., 0.],
[245000., 245000.], [10, 10], [10, 10]])
model, jitter = testSystem_parsing.compute_model(params_arr)
assert model.shape == (4, 2, 2)
params_arr = np.array([1., 0., 0., 0., 0., 245000., 10, 10])
model, jitter = testSystem_parsing.compute_model(params_arr)
assert model.shape == (4, 2)
def test_read_old_orbitize_format():
"""
Test the read_file function when using an old orbitize data file without `quant12_corr` and `instrument` fields.
"""
# Check that main test input is read in with correct values
input_file = os.path.join(orbitize.DATADIR, 'old_orbitize_format.csv')
input_data = read_file(input_file)
# check correlation and instrument are defualts
assert np.isnan(input_data['quant12_corr'][0])
assert input_data['instrument'][0] == 'defsp'
assert np.isnan(input_data['quant12_corr'][1])
assert input_data['instrument'][1] == 'defrd'
assert np.isnan(input_data['quant12_corr'][2])
assert input_data['instrument'][2] == 'defrv'
def test_compute_model():
"""
Test basic functionality of ``System.compute_model()``
"""
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
data_table = read_input.read_file(input_file)
data_table['object'] = 1
testSystem_parsing = system.System(1, data_table, 10., 10.)
params_arr = np.array([[1., 0.5], [0., 0.], [0., 0.], [0., 0.], [0., 0.],
[245000., 245000.], [10, 10], [10, 10]])
model = testSystem_parsing.compute_model(params_arr)
assert model.shape == (4, 2, 2)
params_arr = np.array([1., 0., 0., 0., 0., 245000., 10, 10])
model = testSystem_parsing.compute_model(params_arr)
assert model.shape == (4, 2)
def test_cov_input():
"""
Test including radec and seppa covariances/correlations.
"""
testdir = orbitize.DATADIR
# Check that main test input is read in with correct values
input_file = os.path.join(testdir, 'test_val_cov.csv')
input_data = read_file(input_file)
_compare_table(input_data)
print(input_data)
# Check the covariance column
quant12_corr_truth = [0.25, np.nan, -0.5, np.nan]
assert 'quant12_corr' in input_data.columns
for row, truth in zip(input_data, quant12_corr_truth):
meas = row['quant12_corr']
if np.isnan(truth):
assert np.isnan(meas)
else:
assert truth == pytest.approx(meas)
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))
def test_mcmc_param_idx():
# 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 with fixed mass and plx
n_walkers = 100
myDriver = Driver(input_file,
'MCMC',
1,
1,
0.01,
mcmc_kwargs={
'num_temps': 0,
'num_threads': 1,
'num_walkers': n_walkers
})
# check that sampler.param_idx behaves as expected
assert myDriver.sampler.sampled_param_idx == std_param_idx_fixed_mtot_plx
# construct Driver with no fixed params
myDriver = Driver(input_file,
'MCMC',
1,
1,
0.01,
mass_err=0.1,
plx_err=0.2,
mcmc_kwargs={
'num_temps': 0,
'num_threads': 1,
'num_walkers': n_walkers
})
assert myDriver.sampler.sampled_param_idx == std_param_idx
def test_write_orbitize_input_2():
"""
Test the write_orbitize_input and the read_orbitize_input functions
This test exists with the fail_if_not_removed decorator as a reminder to remove in v2.0
"""
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
test_table = read_file(input_file)
output_file = os.path.join(testdir, 'temp_test_orbitize_input.csv')
# If temp output file already exists, delete it
if os.path.isfile(output_file):
os.remove(output_file)
try: # Catch these tests so that we remove temporary file
# Test that we were able to write the table
write_orbitize_input(test_table,output_file)
assert os.path.isfile(output_file)
# Test that we can read the table and check if it's correct
test_table_2 = read_orbitize_input(output_file)
_compare_table(test_table_2)
finally:
# Remove temporary file
os.remove(output_file)
def test_init_and_add_samples(radec_input=False):
"""
Tests object creation and add_samples() with some simulated posterior samples
Returns results.Results object
"""
if radec_input:
input_file = os.path.join(orbitize.DATADIR, 'test_val_radec.csv')
else:
input_file = os.path.join(orbitize.DATADIR, 'GJ504.csv')
data = read_input.read_file(input_file)
test_system = system.System(1, data, 1, 1)
# Create object
results_obj = results.Results(test_system, sampler_name='testing')
# Simulate some sample draws, assign random likelihoods
n_orbit_draws1 = 1000
sim_post = simulate_orbit_sampling(n_orbit_draws1)
sim_lnlike = np.random.uniform(size=n_orbit_draws1)
# Test adding samples
results_obj.add_samples(sim_post, sim_lnlike) #, labels=std_labels)
# Simulate some more sample draws
n_orbit_draws2 = 2000
sim_post = simulate_orbit_sampling(n_orbit_draws2)
sim_lnlike = np.random.uniform(size=n_orbit_draws2)
# Test adding more samples
results_obj.add_samples(sim_post, sim_lnlike) #, labels=std_labels)
# Check shape of results.post
expected_length = n_orbit_draws1 + n_orbit_draws2
assert results_obj.post.shape == (expected_length, 8)
assert results_obj.lnlike.shape == (expected_length, )
assert results_obj.tau_ref_epoch == 58849
assert results_obj.labels == std_labels
return results_obj
def test_system_setup():
"""
Test that a System object with Hipparcos and Gaia is initialized correctly
"""
hip_num = '027321' # beta Pic
edr3_num = '4792774797545800832'
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(edr3_num, myHip, dr='edr3')
input_file = os.path.join(DATADIR, 'betaPic.csv')
plx = 51.5
num_secondary_bodies = 1
data_table = read_input.read_file(input_file)
betaPic_system = system.System(num_secondary_bodies,
data_table,
1.75,
plx,
hipparcos_IAD=myHip,
gaia=myGaia,
fit_secondary_mass=True,
mass_err=0.01,
plx_err=0.01)
assert betaPic_system.labels == [
'sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'plx', 'pm_ra',
'pm_dec', 'alpha0', 'delta0', 'm1', 'm0'
]
assert betaPic_system.fit_secondary_mass
assert betaPic_system.track_planet_perturbs
print('Go get a coffee. This will take a few mins! :)')
nielsen_iad_refitting_test(IAD_file,
hip_num=hip_num,
saveplot=saveplot,
burn_steps=hip_burn_steps,
mcmc_steps=hip_mcmc_steps)
end = datetime.now()
duration_mins = (end - start).total_seconds() / 60
print(
"Done! This fit took {:.1f} mins on my machine.".format(duration_mins))
### ACTUAL ORBITIZE ###
data_table = read_input.read_file(filename)
print(data_table)
# system parameters
num_secondary_bodies = 1
hipparcos_number = hip_num
hipparcos_filename = IAD_file
# HIP logprob
HD142527_Hip = hipparcos.HipparcosLogProb(hipparcos_filename,
hipparcos_number,
num_secondary_bodies)
# gaia logprob
HD142527_edr3_number = 5994826707951507200
HD142527_Gaia = gaia.GaiaLogProb(HD142527_edr3_number,
HD142527_Hip,
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 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():
"""
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))
