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)
