Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
import matplotlib.pyplot as plt
import time
import orbitize.sampler as sampler
import orbitize.driver
import orbitize.priors as priors
from orbitize.lnlike import chi2_lnlike
from orbitize.kepler import calc_orbit
import orbitize.system
import pdb
testdir = os.path.dirname(os.path.abspath(__file__))
input_file = os.path.join(testdir, 'HD4747.csv')
# perform scale-and-rotate
myDriver = orbitize.driver.Driver(input_file, 'OFTI',
1, 0.84, 53.1836,
mass_err=0.04, plx_err=0.1264,
system_kwargs={'fit_secondary_mass': True, 'tau_ref_epoch': 0}
)
s = myDriver.sampler
samples = s.prepare_samples(10)
if myDriver.system.fit_secondary_mass:
sma, ecc, inc, argp, lan, tau, plx, gamma, sigma, m1, m0 = [samp for samp in samples]
ra, dec, vc = orbitize.kepler.calc_orbit(
s.epochs, sma, ecc, inc, argp, lan, tau, plx, mtot=m1 + m0,
mass_for_Kamp=m0)
v_star = vc*-(m1/m0)
else:
sma, ecc, inc, argp, lan, tau, plx, mtot = [samp for samp in samples]
def test_scale_and_rotate():
# perform scale-and-rotate
myDriver = orbitize.driver.Driver(input_file, 'OFTI',
1, 1.22, 56.95, mass_err=0.08, plx_err=0.26)
s = myDriver.sampler
samples = s.prepare_samples(100)
sma, ecc, inc, argp, lan, tau, plx, mtot = [samp for samp in samples]
ra, dec, vc = orbitize.kepler.calc_orbit(s.epochs, sma, ecc, inc, argp, lan, tau, plx, mtot)
sep, pa = orbitize.system.radec2seppa(ra, dec)
sep_sar, pa_sar = np.median(sep[s.epoch_idx]), np.median(pa[s.epoch_idx])
# test to make sure sep and pa scaled to scale-and-rotate epoch
sar_epoch = s.system.data_table[s.epoch_idx]
assert sep_sar == pytest.approx(sar_epoch['quant1'], abs=sar_epoch['quant1_err'])
assert pa_sar == pytest.approx(sar_epoch['quant2'], abs=sar_epoch['quant2_err'])
def test_scale_and_rotate():
# perform scale-and-rotate
myDriver = orbitize.driver.Driver(input_file, 'OFTI',
1, 1.22, 56.95, mass_err=0.08, plx_err=0.26,
system_kwargs={'fit_secondary_mass': True, 'tau_ref_epoch': 0}
)
s = myDriver.sampler
samples = s.prepare_samples(100)
sma, ecc, inc, argp, lan, tau, plx, gamma, sigma, m1, m0 = [samp for samp in samples]
mtot = m0 + m1
print('samples read')
ra, dec, vc = orbitize.kepler.calc_orbit(s.epochs, sma, ecc, inc, argp, lan, tau, plx, mtot)
sep, pa = orbitize.system.radec2seppa(ra, dec)
sep_sar, pa_sar = np.median(sep[s.epoch_idx]), np.median(pa[s.epoch_idx])
# test to make sure sep and pa scaled to scale-and-rotate epoch
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_fixed_sys_params_sampling():
# test in case of fixed mass and parallax
myDriver = orbitize.driver.Driver(input_file, 'OFTI',
1, 1.22, 56.95)
s = myDriver.sampler
samples = s.prepare_samples(100)
assert np.all(samples[-1] == s.priors[-1])
assert isinstance(samples[-3], np.ndarray)
def test_fixed_sys_params_sampling():
# test in case of fixed mass and parallax
myDriver = orbitize.driver.Driver(input_file, 'OFTI',
1, 1.22, 56.95)
s = myDriver.sampler
samples = s.prepare_samples(100)
assert np.all(samples[-1] == s.priors[-1])
assert isinstance(samples[-3], np.ndarray)