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]