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_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))
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_OFTI_covariances(): """ Test OFTI fits by turning sep/pa measurements to RA/Dec measurements with covariances Needs to be run after test_run_sampler()!! """ # only run if these variables are set. if sma_seppa == 0 or seppa_lnprob_compare is None: print( "Skipping OFTI covariances test because reference data not initalized. Please make sure test_run_sampler is run first." ) return # read in seppa data table and turn into raddec data table data_table = orbitize.read_input.read_file(input_file) data_ra, data_dec = system.seppa2radec(data_table['quant1'], data_table['quant2']) data_raerr, data_decerr, data_radeccorr = [], [], [] for row in data_table: raerr, decerr, radec_corr = system.transform_errors(row['quant1'], row['quant2'], row['quant1_err'], row['quant2_err'], 0, system.seppa2radec, nsamps=10000000) data_raerr.append(raerr) data_decerr.append(decerr) data_radeccorr.append(radec_corr) data_table['quant1'] = data_ra data_table['quant2'] = data_dec data_table['quant1_err'] = np.array(data_raerr) data_table['quant2_err'] = np.array(data_decerr) data_table['quant12_corr'] = np.array(data_radeccorr) data_table['quant_type'] = np.array(['radec' for _ in data_table]) # initialize system my_sys = system.System(1, data_table, 1.22, 56.95, mass_err=0.08, plx_err=0.26) # initialize sampler s = sampler.OFTI(my_sys) # change eccentricity prior my_sys.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 orbits = s.run_sampler(1000, num_cores=4) # test that lnlikes being saved are correct returned_lnlike_test = s.results.lnlike[0] computed_lnlike_test = s._logl(orbits[0]) assert returned_lnlike_test == pytest.approx(computed_lnlike_test, abs=0.01) # test that the lnlike is very similar to the values computed in seppa space ref_params, ref_lnlike = seppa_lnprob_compare computed_lnlike_ref = s._logl(ref_params) assert ref_lnlike == pytest.approx( computed_lnlike_ref, abs=0.05) # 5% differencesin lnprob is allowable. idx = s.system.param_idx sma = np.median([x[idx['sma1']] for x in orbits]) ecc = np.median([x[idx['ecc1']] for x in orbits]) inc = np.median([x[idx['inc1']] for x in orbits]) # test against seppa fits to see they are similar assert sma_seppa == pytest.approx(sma, abs=0.2 * sma_seppa)
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_2planet_nomass(): """ Compare multiplanet to rebound for planets with mass. """ # generate a planet orbit mjup = u.Mjup.to(u.Msun) mass_b = 0 * mjup mass_c = 0 * mjup params = np.array([ 10, 0.1, np.radians(45), np.radians(45), np.radians(45), 0.5, 3, 0.1, np.radians(45), np.radians(190), np.radians(45), 0.2, 1, mass_b, mass_c, 1.5 - mass_b - mass_c ]) tau_ref_epoch = 0 epochs = np.linspace(0, 365.25 * 4, 100) + tau_ref_epoch # nearly the full period, MJD # doesn't matter that this is right, just needs to be the same size. below doesn't include effect of c # just want to generate some measurements of plaent b to test compute model b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit( epochs, params[0], params[1], params[2], params[3], params[4], params[5], params[-2], params[-1], tau_ref_epoch=tau_ref_epoch) # generate some fake measurements of planet b, just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), b_ra_model, np.zeros(b_ra_model.shape), b_dec_model, np.zeros(b_dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_2planet.csv") t.write(filename) # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, params[-1], params[-2], tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # generate measurement radec_orbitize, _ = sys.compute_model(params) b_ra_orb = radec_orbitize[:, 0] b_dec_orb = radec_orbitize[:, 1] # now project the orbit with rebound b_manom = basis.tau_to_manom(epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch) c_manom = basis.tau_to_manom(epochs[0], params[0 + 6], params[-1] + params[-2], params[5 + 6], tau_ref_epoch) sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=params[-1]) # add two planets sim.add(m=mass_c, a=params[0 + 6], e=params[1 + 6], M=c_manom, omega=params[3 + 6], Omega=params[4 + 6] + np.pi / 2, inc=params[2 + 6]) sim.add(m=mass_b, a=params[0], e=params[1], M=b_manom, omega=params[3], Omega=params[4] + np.pi / 2, inc=params[2]) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation b_ra_reb = [] b_dec_reb = [] for t in epochs: sim.integrate(t / 365.25) b_ra_reb.append(-(ps[2].x - ps[0].x)) # ra is negative x b_dec_reb.append(ps[2].y - ps[0].y) b_ra_reb = np.array(b_ra_reb) b_dec_reb = np.array(b_dec_reb) diff_ra = b_ra_reb - b_ra_orb / params[6 * 2] diff_dec = b_dec_reb - b_dec_orb / params[6 * 2] # should be as good as the one planet case assert np.all(np.abs(diff_ra) / (params[0] * params[6 * 2]) < 1e-9) assert np.all(np.abs(diff_dec) / (params[0] * params[6 * 2]) < 1e-9)
HD142527_Hip, dr='edr3') # more system parameters m0 = 2.05 # [Msol] plx = 6.35606723729484 # [mas] fit_secondary_mass = True mass_err = 0.5 # [Msol] plx_err = 0.04714455423 # [mas] HD142527_system = system.System(num_secondary_bodies, data_table, m0, plx, hipparcos_IAD=HD142527_Hip, gaia=HD142527_Gaia, fit_secondary_mass=fit_secondary_mass, mass_err=mass_err, plx_err=plx_err) # 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
def test_2planet_massive(): """ Compare multiplanet to rebound for planets with mass. """ # generate a planet orbit mjup = u.Mjup.to(u.Msun) mass_b = 12 * mjup mass_c = 9 * mjup params = np.array([ 10, 0.1, np.radians(45), np.radians(45), np.radians(45), 0.5, 3, 0.1, np.radians(45), np.radians(190), np.radians(45), 0.2, 50, mass_b, mass_c, 1.5 - mass_b - mass_c ]) params_noc = np.array([ 10, 0.1, np.radians(45), np.radians(45), np.radians(45), 0.5, 3, 0.1, np.radians(45), np.radians(190), np.radians(45), 0.2, 50, mass_b, 0, 1.5 - mass_b ]) tau_ref_epoch = 0 epochs = np.linspace(0, 365.25 * 10, 100) + tau_ref_epoch # nearly the full period, MJD # doesn't matter that this is right, just needs to be the same size. below doesn't include effect of c # just want to generate some measurements of plaent b to test compute model b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit( epochs, params[0], params[1], params[2], params[3], params[4], params[5], params[-2], params[-1], tau_ref_epoch=tau_ref_epoch) # generate some fake measurements of planet b, just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), b_ra_model, np.zeros(b_ra_model.shape), b_dec_model, np.zeros(b_dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_2planet_outer.csv") t.write(filename) #### TEST THE OUTER PLANET #### # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, params[-1], params[-4], tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # generate measurement radec_orbitize, _ = sys.compute_model(params) b_ra_orb = radec_orbitize[:, 0] b_dec_orb = radec_orbitize[:, 1] # debug, generate measurement without c having any mass radec_orb_noc, _ = sys.compute_model(params_noc) b_ra_orb_noc = radec_orb_noc[:, 0] b_dec_orb_noc = radec_orb_noc[:, 1] # check that planet c's perturbation is imprinted (nonzero)) #assert np.all(b_ra_orb_noc != b_ra_orb) # now project the orbit with rebound b_manom = basis.tau_to_manom(epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch) c_manom = basis.tau_to_manom(epochs[0], params[0 + 6], params[-1] + params[-2], params[5 + 6], tau_ref_epoch) sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=params[-1]) # add two planets sim.add(m=mass_c, a=params[0 + 6], e=params[1 + 6], M=c_manom, omega=params[3 + 6], Omega=params[4 + 6] + np.pi / 2, inc=params[2 + 6]) sim.add(m=mass_b, a=params[0], e=params[1], M=b_manom, omega=params[3], Omega=params[4] + np.pi / 2, inc=params[2]) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation b_ra_reb = [] b_dec_reb = [] for t in epochs: sim.integrate(t / 365.25) b_ra_reb.append(-(ps[2].x - ps[0].x)) # ra is negative x b_dec_reb.append(ps[2].y - ps[0].y) b_ra_reb = np.array(b_ra_reb) b_dec_reb = np.array(b_dec_reb) diff_ra = b_ra_reb - b_ra_orb / params[6 * 2] diff_dec = b_dec_reb - b_dec_orb / params[6 * 2] # we placed the planets far apart to minimize secular interactions but there are still some, so relax precision assert np.all(np.abs(diff_ra) / (params[0]) < 1e-3) assert np.all(np.abs(diff_dec) / (params[0]) < 1e-3) ###### NOW TEST THE INNER PLANET ####### # generate some fake measurements of planet c, just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int) * 2, b_ra_model, np.zeros(b_ra_model.shape), b_dec_model, np.zeros(b_dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_2planet_inner.csv") t.write(filename) # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, params[-1], params[-2], tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # generate measurement radec_orbitize, _ = sys.compute_model(params) c_ra_orb = radec_orbitize[:, 0] c_dec_orb = radec_orbitize[:, 1] # start the REBOUND sim again sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=params[-1]) # add two planets sim.add(m=mass_c, a=params[0 + 6], e=params[1 + 6], M=c_manom, omega=params[3 + 6], Omega=params[4 + 6] + np.pi / 2, inc=params[2 + 6]) sim.add(m=mass_b, a=params[0], e=params[1], M=b_manom, omega=params[3], Omega=params[4] + np.pi / 2, inc=params[2]) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation c_ra_reb = [] c_dec_reb = [] for t in epochs: sim.integrate(t / 365.25) c_ra_reb.append(-(ps[1].x - ps[0].x)) # ra is negative x c_dec_reb.append(ps[1].y - ps[0].y) c_ra_reb = np.array(c_ra_reb) c_dec_reb = np.array(c_dec_reb) diff_ra = c_ra_reb - c_ra_orb / params[6 * 2] diff_dec = c_dec_reb - c_dec_orb / params[6 * 2] # planet is 3 times closer, so roughly 3 times larger secular errors. assert np.all(np.abs(diff_ra) / (params[0]) < 3e-3) assert np.all(np.abs(diff_dec) / (params[0]) < 3e-3)
def test_1planet(): """ Sanity check that things agree for 1 planet case """ # 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 mtot = 1 tau_ref_epoch = 0 mjup = u.Mjup.to(u.Msun) mass_b = 12 * mjup epochs = np.linspace(0, 300, 100) + tau_ref_epoch # nearly the full period, MJD ra_model, dec_model, vz_model = kepler.calc_orbit( epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch) # generate some fake measurements just 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 the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(1, astrom_dat, mtot, plx, tau_ref_epoch=tau_ref_epoch) params = np.array([sma, ecc, inc, aop, pan, tau, plx, mtot]) radec_orbitize, _ = sys.compute_model(params) ra_orb = radec_orbitize[:, 0] dec_orb = radec_orbitize[:, 1] # now project the orbit with rebound manom = basis.tau_to_manom(epochs[0], sma, mtot, tau, tau_ref_epoch) sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=mtot - mass_b) # add one planet sim.add(m=mass_b, a=sma, e=ecc, M=manom, omega=aop, Omega=pan + np.pi / 2, inc=inc) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation ra_reb = [] dec_reb = [] for t in epochs: sim.integrate(t / 365.25) ra_reb.append(-(ps[1].x - ps[0].x)) # ra is negative x dec_reb.append(ps[1].y - ps[0].y) ra_reb = np.array(ra_reb) dec_reb = np.array(dec_reb) diff_ra = ra_reb - ra_orb / plx diff_dec = dec_reb - dec_orb / plx assert np.all(np.abs(diff_ra) < 1e-9) assert np.all(np.abs(diff_dec) < 1e-9)
n_temps = 20 n_walkers = 1000 n_threads = mp.cpu_count() total_orbits = 10000000 # n_walkers x num_steps_per_walker burn_steps = 50000 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")
betaPic_Hip = HipparcosLogProb(hipparcos_filename, hipparcos_number, num_secondary_bodies) betaPic_gaia = GaiaLogProb(gaia_dr2_number, betaPic_Hip, dr='dr2') # betaPic_gaia = GaiaLogProb( # gaia_edr3_number, betaPic_Hip, dr='edr3' # ) else: fit_secondary_mass = False betaPic_Hip = None betaPic_gaia = None betaPic_system = system.System(num_secondary_bodies, data_table, 1.75, plx, hipparcos_IAD=betaPic_Hip, gaia=betaPic_gaia, fit_secondary_mass=fit_secondary_mass, mass_err=0.01, plx_err=0.01) m0_or_mtot_prior = priors.UniformPrior(1.5, 2.0) # set uniform parallax prior plx_index = betaPic_system.param_idx['plx'] betaPic_system.sys_priors[plx_index] = priors.UniformPrior( plx - 1.0, plx + 1.0) # set prior on Omega, since we know that know direction of orbital motion from RV pan_index = betaPic_system.param_idx['pan1'] betaPic_system.sys_priors[pan_index] = priors.UniformPrior(0, np.pi)
def test_hipparcos_api(): """ Check that error is caught for a star with solution type != 5 param, and that doing an RV + Hipparcos IAD fit produces the expected list of Prior objects. """ # check sol type != 5 error message hip_num = '000025' num_secondary_bodies = 1 path_to_iad_file = '{}H{}.d'.format(DATADIR, hip_num) try: _ = HipparcosLogProb(path_to_iad_file, hip_num, num_secondary_bodies) assert False, 'Test failed.' except ValueError: pass # check that RV + Hip gives correct prior array labels hip_num = '027321' # beta Pic num_secondary_bodies = 1 path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num) myHip = HipparcosLogProb(path_to_iad_file, hip_num, num_secondary_bodies) input_file = os.path.join(DATADIR, 'HD4747.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, fit_secondary_mass=True) # test that `fit_secondary_mass` and `track_planet_perturbs` keywords are set appropriately assert mySys.fit_secondary_mass assert mySys.track_planet_perturbs assert len(mySys.sys_priors) == 15 # 7 orbital params + 2 mass params + # 4 Hip nuisance params + # 2 RV nuisance params assert mySys.labels == [ 'sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'plx', 'pm_ra', 'pm_dec', 'alpha0', 'delta0', 'gamma_defrv', 'sigma_defrv', 'm1', 'm0' ] # test that `fit_secondary_mass` and `track_planet_perturbs` keywords are # set appropriately for non-Hipparcos system noHip_system = system.System(num_secondary_bodies, data_table_with_rvs, 1.0, 1.0, hipparcos_IAD=None, fit_secondary_mass=False, mass_err=0.01, plx_err=0.01) assert not noHip_system.fit_secondary_mass assert not noHip_system.track_planet_perturbs # check that negative residuals are rejected properly hip_num = '000026' # contains one negative residual num_secondary_bodies = 1 path_to_iad_file = '{}H{}.d'.format(DATADIR, hip_num) raw_iad_data = np.transpose(np.loadtxt(path_to_iad_file)) rejected_scansHip = HipparcosLogProb(path_to_iad_file, hip_num, num_secondary_bodies) assert len(rejected_scansHip.cos_phi) == len(raw_iad_data[0]) - 1
def test_period_basis(): """ For both MCMC and OFTI formats, make the conversion to standard basis and go back to original basis and check to see original params are retrieved. Do this with system mass parameter, single companion, and two companions. """ # 1. With System Total Mass filename = "{}/GJ504.csv".format(DATADIR) data_table = read_input.read_file(filename) my_system = system.System( 1, data_table, 1.75, 51.44, mass_err=0.05, plx_err=0.12, fitting_basis='Period' ) 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_period_basis(conversion) assert np.allclose(original, sample_copy[:, 0]) # OFTI Format conversions = my_system.basis.to_standard_basis(samples) original = my_system.basis.to_period_basis(conversions) assert np.allclose(original, sample_copy) # 2. 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='Period' ) 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_period_basis(conversion) assert np.allclose(original, sample_copy[:, 0]) # 3. 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='Period' ) 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_period_basis(conversion) assert np.allclose(original, sample_copy[:, 0]) # OFTI Format conversions = my_system.basis.to_standard_basis(samples) original = my_system.basis.to_period_basis(conversions) assert np.allclose(original, sample_copy)
def test_xyz_basis(): """ For both MCMC and OFTI param formats, make the conversion to the standard basis from XYZ basis and back to verify the valdity of conversions. Do this with a single companion and with two companions. """ # 1. Single Body filename = '{}/xyz_test_data.csv'.format(DATADIR) data = read_input.read_file(filename) single = data[np.where(data['object'] == 1)[0]] my_system = system.System( 1, single, 1.22, 56.89, mass_err=0.05, plx_err=0.12, fitting_basis='XYZ' ) num_samples = 1000 # Do more samples to be safe 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() conversion = my_system.basis.to_standard_basis(samples) locs = np.where((conversion[1, :] >= 1.0) | (conversion[1, :] < 0.))[0] sample_copy = np.delete(sample_copy, locs, axis=1) # Test MCMC test = sample_copy[:, 0].copy() conversion = my_system.basis.to_standard_basis(test) original = my_system.basis.to_xyz_basis(conversion) assert np.allclose(original, sample_copy[:, 0]) # Test OFTI conversions = my_system.basis.to_standard_basis(sample_copy.copy()) original = my_system.basis.to_xyz_basis(conversions) assert np.allclose(original, sample_copy) # 2. Multi Body my_system = system.System( 2, data, 1.22, 56.89, mass_err=0.05, plx_err=0.12, fitting_basis='XYZ' ) num_samples = 1000 # Do more samples to be safe 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() conversion = my_system.basis.to_standard_basis(samples) locs = np.where( (conversion[[1, 7], :] >= 1.0) | (conversion[[1, 7], :] < 0.) )[1] locs = np.unique(locs) sample_copy = np.delete(sample_copy, locs, axis=1) # Test MCMC test = sample_copy[:, 0].copy() conversion = my_system.basis.to_standard_basis(test) original = my_system.basis.to_xyz_basis(conversion) assert np.allclose(original, sample_copy[:, 0]) # Test OFTI conversions = my_system.basis.to_standard_basis(sample_copy.copy()) original = my_system.basis.to_xyz_basis(conversions) assert np.allclose(original, sample_copy)