Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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]
_compare_table(myDriver.system.data_table)
def test_create_driver_from_filename():
"""
Test creation of Driver object from filename as input
"""
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
myDriver = driver.Driver(input_file, # path to data file
'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_run_sampler():
# 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 eccentricity prior
myDriver.system.sys_priors[1] = priors.LinearPrior(-2.18, 2.01)
# test num_samples=1
s.run_sampler(0, num_samples=1)
# test to make sure outputs are reasonable
start = time.time()
orbits = s.run_sampler(1000, num_cores=4)
end = time.time()
print()
# expected values from Blunt et al. (2017)
sma_exp = 48.
ecc_exp = 0.19
inc_exp = np.radians(140)
# test to make sure OFTI values are within 20% of expectations
assert sma == pytest.approx(sma_exp, abs=0.2*sma_exp)
assert ecc == pytest.approx(ecc_exp, abs=0.2*ecc_exp)
assert inc == pytest.approx(inc_exp, abs=0.2*inc_exp)
# test with only one core
orbits = s.run_sampler(100, num_cores=1)
# test with only one epoch
myDriver = orbitize.driver.Driver(input_file_1epoch, 'OFTI',
1, 1.22, 56.95, mass_err=0.08, plx_err=0.26)
s = myDriver.sampler
s.run_sampler(1)
print()
def test_pt_mcmc_runs(num_threads=1):
"""
Tests the PTMCMC sampler by making sure it even runs
"""
# use the test_csv dir
input_file = os.path.join(orbitize.DATADIR, 'test_val_rv.csv')
myDriver = Driver(input_file, 'MCMC', 1, 1, 0.01,
# mass_err=0.05, plx_err=0.01,
system_kwargs={'fit_secondary_mass': True, 'tau_ref_epoch': 0},
mcmc_kwargs={'num_temps': 2, 'num_threads': num_threads, 'num_walkers': 100}
)
# run it a little (tests 0 burn-in steps)
myDriver.sampler.run_sampler(100)
# run it a little more
myDriver.sampler.run_sampler(1000, burn_steps=1)
# run it a little more (tests adding to results object)
myDriver.sampler.run_sampler(1000, burn_steps=1)
fixed_index = [index for index, fixed_param in myDriver.sampler.fixed_params]
clean_params = np.delete(myDriver.sampler.results.post[0], fixed_index)