def test_with_rv():
'''
Make sure the labels are generated correctly for all of the basis sets for the configuration of
the driver where RV data is supplied. Again, for the XYZ basis, expect there to be errors thrown.
'''
filename = "{}/HD4747.csv".format(DATADIR)
num_secondary_bodies = 1
system_mass = 0.84 # [Msol]
plx = 53.18 # [mas]
mass_err = 0.04 # [Msol]
plx_err = 0.12 # [mas]
# (1) Single Body
for basis in basis_names:
if basis != 'SemiAmp':
expected_labels = [item + '1' for item in basis_names[basis]] + [
'plx', 'gamma_defrv', 'sigma_defrv', 'm1', 'm0'
]
else:
expected_labels = [item + '1' for item in basis_names[basis]
] + ['plx', 'gamma_defrv', 'sigma_defrv', 'm0']
if basis == 'XYZ':
with pytest.raises(Exception) as excinfo:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0,
'fitting_basis': basis
})
assert str(
excinfo.value
) == "For now, the XYZ basis requires data in RA and DEC offsets."
else:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0,
'fitting_basis': basis
})
assert expected_labels == my_driver.system.labels
def test_rv_multiple_inst():
filename = '{}/HR7672_joint.csv'.format(DATADIR)
num_secondary_bodies = 1
system_mass = 1.08 # [Msol]
plx = 56.2 # [mas]
mass_err = 0.04 # [Msol]
plx_err = 0.01 # [mas]
# MCMC parameters
num_temps = 5
num_walkers = 30
num_threads = 2
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0
},
mcmc_kwargs={
'num_temps': num_temps,
'num_walkers': num_walkers,
'num_threads': num_threads
})
total_orbits = 500
burn_steps = 10
thin = 2
m = my_driver.sampler
m.run_sampler(total_orbits, burn_steps=burn_steps, thin=thin)
epochs = my_driver.system.data_table['epoch']
orbit_plot_fig = m.results.plot_orbits(
object_to_plot=
1, # Plot orbits for the first (and only, in this case) companion
num_orbits_to_plot=
1, # Will plot 100 randomly selected orbits of this companion
start_mjd=epochs[
0], # Minimum MJD for colorbar (here we choose first data epoch)
rv_time_series=True,
plot_astrometry_insts=True)
# Test plotting with multiple orbits
orbit_plot_fig = m.results.plot_orbits(
object_to_plot=
1, # Plot orbits for the first (and only, in this case) companion
num_orbits_to_plot=
10, # Will plot 100 randomly selected orbits of this companion
start_mjd=epochs[
0], # Minimum MJD for colorbar (here we choose first data epoch)
rv_time_series=True,
plot_astrometry_insts=True)
def test_create_driver_from_filename():
"""
Test creation of Driver object from filename as input
"""
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'test_val.csv')
myDriver = driver.Driver(input_file, # path to data file
'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_create_driver_from_filename():
"""
Test creation of Driver object from filename as input
"""
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
myDriver = driver.Driver(
input_file, # path to data file
'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]
system_kwargs={'fit_secondary_mass': True})
_compare_table(myDriver.system.data_table)
def test_system_kwargs():
"""
Test additional arguments to the system class
"""
input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
myDriver = driver.Driver(
input_file, # path to data file
'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,
system_kwargs={"tau_ref_epoch": 50000}) # parallax error [mas]
assert myDriver.system.tau_ref_epoch == 50000
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]
_compare_table(myDriver.system.data_table)
def test_chop_chains():
'''
First run MCMC sampler to generate results object and make a call to 'chop_chains'
function afterwards.
'''
filename = "{}/HD4747.csv".format(DATADIR)
num_secondary_bodies = 1
system_mass = 0.84
plx = 53.18
mass_err = 0.04
plx_err = 0.12
num_temps = 5
num_walkers = 40
num_threads = mp.cpu_count()
total_orbits = 5000
burn_steps = 10
thin = 2
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0
},
mcmc_kwargs={
'num_temps': num_temps,
'num_walkers': num_walkers,
'num_threads': num_threads
})
my_driver.sampler.run_sampler(total_orbits,
burn_steps=burn_steps,
thin=thin)
my_driver.sampler.chop_chains(burn=25, trim=25)
mcmc_sys = my_driver.system
mcmc_result = my_driver.sampler.results
verify_results_data(mcmc_result, mcmc_sys)
python -m cProfile -o profiler_output.txt profile_orbitize.py
The `-m cProfile` will attach the cProfile profiler to this script when it runs
The `-o profiler_output.txt` will save its output to "profiler_output.txt".
If you have any issues running the above command, your python configuration might be different than what
we are assuing here. Ask a TA to help you figure out what needs to be changed if needed!
Note on parallelized code:
orbitize! by default is parallelized, but we will avoid doing that here.
Using parallelism will make the cprofiler output more confusing than necessary. We just want to know what takes the most time,
summed up across multiple processes. Doing this on parallelized code, each process must save its own cProfile output and you
need to combine the outputs afterwards (using `pstats.add()`). It is bset to avoid this when possible. Here, we do a very short
orbit fit (probably not good enough for science) to benchmark our code on a single core in a short amount of time.
"""
import orbitize
from orbitize import driver
myDriver = driver.Driver(
'{}/GJ504.csv'.format(orbitize.DATADIR), # data file
'OFTI', # choose from: ['OFTI', 'MCMC']
1, # number of planets in system
1.22, # total system mass [M_sun]
56.95, # system parallax [mas]
mass_err=0.08, # mass error [M_sun]
plx_err=0.26 # parallax error [mas]
)
orbits = myDriver.sampler.run_sampler(1000, num_cores=1)
if "Jason" in suffix:
plx = 24.76 #Jason
plx_err = 0.64 #Jason
print(plx, plx_err)
if 1:
out_pngs = os.path.join(astrometry_DATADIR, "figures")
my_driver = driver.Driver(
filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
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,
def test_with_hip_iad():
'''
Make sure the labels are generated correctly for all of the basis sets for the configuration of
the driver where Hipparcos data is supplied. Again, for the XYZ basis, expect there to be errors thrown.
'''
filename = "{}/betaPic.csv".format(DATADIR)
hip_num = '027321'
hipp_filename = "{}/HIP{}.d".format(DATADIR, hip_num)
num_secondary_bodies = 1
system_mass = 1.75
mass_err = 0
plx = 51.44
plx_err = 0.12
myHip = hipparcos.HipparcosLogProb(hipp_filename, hip_num,
num_secondary_bodies)
hip_labels = ['pm_ra', 'pm_dec', 'alpha0', 'delta0']
for basis in basis_names:
if basis != 'SemiAmp':
expected_labels = [item + '1' for item in basis_names[basis]
] + ['plx'] + hip_labels + ['m1', 'm0']
else:
expected_labels = [item + '1' for item in basis_names[basis]
] + ['plx'] + hip_labels + ['m0']
if basis == 'XYZ': # Should throw an error for XYZ basis
with pytest.raises(Exception) as excinfo:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0,
'hipparcos_IAD': myHip,
'fitting_basis': basis
})
assert str(
excinfo.value
) == "For now, the XYZ basis requires data in RA and DEC offsets."
else:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0,
'hipparcos_IAD': myHip,
'fitting_basis': basis
})
assert expected_labels == my_driver.system.labels
def test_no_extra_data():
"""
Make sure that the labels are generated properly for all of the basis sets
for configurations of the driver where (1) system mass, (2) single companion
mass, and (3) two companion masses are being fitted. In any case, RV or
Hipparcos parameters are not being fitted.
For XYZ, expect there to be exceptions thrown when making the driver.
"""
filename = "{}/GJ504.csv".format(DATADIR)
num_secondary_bodies = 1
system_mass = 1.75 # [Msol]
plx = 51.44 # [mas]
mass_err = 0.05 # [Msol]
plx_err = 0.12 # [mas]
# (1) System Mass (fit secondary mass is false)
for basis in basis_names:
if basis != 'SemiAmp':
expected_labels = [item + '1' for item in basis_names[basis]
] + ['plx', 'mtot']
else:
expected_labels = [item + '1'
for item in basis_names[basis]] + ['plx', 'm0']
if basis == 'XYZ': # Should throw error for XYZ basis
with pytest.raises(Exception) as excinfo:
my_driver = driver.Driver(
filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={'fitting_basis': basis})
assert str(
excinfo.value
) == "For now, the XYZ basis requires data in RA and DEC offsets."
else:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={'fitting_basis': basis})
assert expected_labels == my_driver.system.labels
# (2) Single Companion
for basis in basis_names:
if basis != 'SemiAmp':
expected_labels = [item + '1' for item in basis_names[basis]
] + ['plx', 'm1', 'm0']
else:
expected_labels = [item + '1'
for item in basis_names[basis]] + ['plx', 'm0']
if basis == 'XYZ': # Should throw error for XYZ basis
with pytest.raises(Exception) as excinfo:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0,
'fitting_basis': basis
})
assert str(
excinfo.value
) == "For now, the XYZ basis requires data in RA and DEC offsets."
else:
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0,
'fitting_basis': basis
})
assert expected_labels == my_driver.system.labels
# (3) Multiple Companions
filename = "{}/test_val_multi.csv".format(DATADIR)
for basis in basis_names:
init_labels = [item + '1' for item in basis_names[basis]
] + [item + '2' for item in basis_names[basis]]
if basis != 'SemiAmp':
expected_labels = init_labels + ['plx', 'm1', 'm2', 'm0']
else:
expected_labels = init_labels + ['plx', 'm0']
if basis == 'XYZ': # Should throw error for XYZ basis
with pytest.raises(Exception) as excinfo:
my_driver = driver.Driver(filename,
'MCMC',
2,
1.52,
24.76,
mass_err=0.15,
plx_err=0.64,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': True,
'fitting_basis': basis
})
assert str(
excinfo.value
) == "For now, the epoch with the lowest sepparation error should not be one of the last two entries for body1"
else:
my_driver = driver.Driver(filename,
'MCMC',
2,
1.52,
24.76,
mass_err=0.15,
plx_err=0.64,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': True,
'fitting_basis': basis
})
assert expected_labels == my_driver.system.labels
def test_rv_default_inst():
# Initialize Driver to Run MCMC
filename = '{}/HD4747.csv'.format(DATADIR)
num_secondary_bodies = 1
system_mass = 0.84 # [Msol]
plx = 53.18 # [mas]
mass_err = 0.04 # [Msol]
plx_err = 0.12 # [mas]
num_temps = 5
num_walkers = 30
num_threads = mp.cpu_count() # or a different number if you prefer
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': True,
'tau_ref_epoch': 0
},
mcmc_kwargs={
'num_temps': num_temps,
'num_walkers': num_walkers,
'num_threads': num_threads
})
total_orbits = 1000
burn_steps = 10
thin = 2
# Run Quick Sampler
m = my_driver.sampler
m.run_sampler(total_orbits, burn_steps=burn_steps, thin=thin)
epochs = my_driver.system.data_table['epoch']
# Test plotting with single orbit
orbit_plot_fig = m.results.plot_orbits(
object_to_plot=
1, # Plot orbits for the first (and only, in this case) companion
num_orbits_to_plot=
1, # Will plot 100 randomly selected orbits of this companion
start_mjd=epochs[
3], # Minimum MJD for colorbar (here we choose first data epoch)
rv_time_series=True,
plot_astrometry_insts=True)
# Test plotting with multiple orbits
orbit_plot_fig = m.results.plot_orbits(
object_to_plot=
1, # Plot orbits for the first (and only, in this case) companion
num_orbits_to_plot=
10, # Will plot 100 randomly selected orbits of this companion
start_mjd=epochs[
3], # Minimum MJD for colorbar (here we choose first data epoch)
rv_time_series=True,
plot_astrometry_insts=True)
plx_err = 0.042361263 # [mas]
# MCMC parameters
num_temps = 10
num_walkers = 50
num_threads = int(mp.cpu_count() * 2 / 3)
# init the driver class that will run orbitize
my_driver = driver.Driver(filename,
'MCMC',
num_secondary_bodies,
system_mass,
plx,
mass_err=mass_err,
plx_err=plx_err,
system_kwargs={
'fit_secondary_mass': False,
'tau_ref_epoch': 0
},
mcmc_kwargs={
'num_temps': num_temps,
'num_walkers': num_walkers,
'num_threads': num_threads
})
print('Driver Initialized...')
# assign the sampler
m = my_driver.sampler
# getting the system object:
sys = my_driver.system
# get parameter labels
lab = sys.param_idx
