def test_read_formatted_file():
"""
Test the read_formatted_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_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(testdir, 'test_val_radec.csv')
read_formatted_file(input_file_radec)
def test_add_and_clear_results():
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_formatted_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_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 test_custom_likelihood():
"""
Tests the inclusion of a custom likelihood function in the code
"""
# use the test_csv dir
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'GJ504.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 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_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 = 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
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', 'epp1', 'sma2', 'ecc2', 'inc2',
'aop2', 'pan2', 'epp2', 'plx', 'mtot'
]
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_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, 'test_val.csv')
data_table = read_input.read_formatted_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': 2,
'num_threads': num_threads,
'num_walkers': n_walkers
})
# 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)
# 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)
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_formatted_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_write_read_orbitize_input():
"""
Test the write_orbitize_input and the read_orbitize_input functions
"""
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
test_table = read_formatted_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_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
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, 'test_val.csv')
data_table = read_input.read_formatted_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]
