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_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 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 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_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))
# set uniform primary mass prior m0_index = HD142527_system.param_idx['m0'] HD142527_system.sys_priors[m0_index] = priors.GaussianPrior(2.05, 0.3) # set uniform primary mass prior m1_index = HD142527_system.param_idx['m1'] HD142527_system.sys_priors[m1_index] = priors.GaussianPrior(0.25, 0.2) # MCMC parameters num_temps = 20 num_walkers = 50 num_threads = int(mp.cpu_count() * 2 / 3) # init driver HD142527_sampler = sampler.MCMC(HD142527_system, num_threads=num_threads, num_temps=num_temps, num_walkers=num_walkers) print('Driver Initialized! Running MCMC...') HD142527_sampler.run_sampler(total_orbits, burn_steps=burn_steps, thin=thin) myResults = HD142527_sampler.results # save posterior to disk savefile = 'paperdraft_posterior.hdf5' myResults.save_results(savefile) print('Saved Posterior!') starttime = Time(datetime.strptime('1990 January 1',
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_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
sysrv=sysrv, sysrv_err=sysrv_err, mcmc_kwargs={ 'num_temps': num_temps, 'num_walkers': num_walkers, 'num_threads': num_threads }, system_kwargs={"restrict_angle_ranges": restrict_angle_ranges}, ) if len(planet) == 1: my_driver.system.sys_priors[0] = priors.JeffreysPrior(1, 1e2) my_driver.sampler = sampler.MCMC(my_driver.system, num_temps=num_temps, num_walkers=num_walkers, num_threads=num_threads, like='chi2_lnlike', custom_lnlike=None) if 1 and len(planet) == 1: _filename = os.path.join( astrometry_DATADIR, "figures", "HR_8799_" + planet, 'chain_' + rv_str + '_' + planet + '_it{0}_{1}_{2}_{3}_{4}_{5}_coplanar.fits'.format( itnum - 1, num_temps, num_walkers, 100000, thin, uservs)) print(_filename) with pyfits.open(_filename) as hdulist: #total_orbits//num_walkers print(hdulist[0].data.shape) my_driver.sampler.curr_pos = hdulist[0].data[:, :, -1, :] my_driver.sampler.save_intermediate = os.path.join(
tau_ref_epoch = 50000 # Read in data data_table = read_input.read_file(datafile) # Initialize System object which stores data & sets priors my_system = system.System(num_secondary_bodies, data_table, system_mass, plx, mass_err=mass_err, plx_err=plx_err, tau_ref_epoch=tau_ref_epoch) my_sampler = sampler.MCMC(my_system, n_temps, n_walkers, n_threads) # Run the sampler to compute some orbits, yeah! # Results stored in bP_sampler.chain and bP_sampler.lnlikes my_sampler.run_sampler(total_orbits, burn_steps=burn_steps, thin=10) my_sampler.results.save_results("hr8799e_gravity_chains.hdf5") #import orbitize.results as results #my_sampler.results = results.Results() #my_sampler.results.load_results("hr8799e_gravity_chains.hdf5") # make corner plot fig = my_sampler.results.plot_corner() plt.savefig('corner_hr8799e_gravity.png', dpi=250)
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, 'GJ504.csv') data_table = read_input.read_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]
assert not betaPic_system.track_planet_perturbs # set uniform mtot prior mtot_index = betaPic_system.param_idx['mtot'] betaPic_system.sys_priors[mtot_index] = m0_or_mtot_prior # run MCMC num_threads = 100 num_temps = 20 num_walkers = 1000 num_steps = 1000000 # n_walkers x n_steps_per_walker burn_steps = 1000 thin = 100 betaPic_sampler = sampler.MCMC(betaPic_system, num_threads=num_threads, num_temps=num_temps, num_walkers=num_walkers) betaPic_sampler.run_sampler(num_steps, burn_steps=burn_steps, thin=thin) # save chains betaPic_sampler.results.save_results('{}/betaPic_IAD{}.hdf5'.format( savedir, fit_IAD)) # make corner plot fig = betaPic_sampler.results.plot_corner() plt.savefig('{}/corner_IAD{}.png'.format(savedir, fit_IAD), dpi=250) # make orbit plot fig = betaPic_sampler.results.plot_orbits() plt.savefig('{}/orbit_IAD{}.png'.format(savedir, fit_IAD), dpi=250)