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))