Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
# 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
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'])
print('done asserting')
# test scale-and-rotate for orbits run all the way through OFTI
s.run_sampler(100)
# test orbit plot generation
s.results.plot_orbits(start_mjd=s.epochs[0])
samples = s.results.post
sma = samples[:, 0]
ecc = samples[:, 1]
def test_radec2seppa():
"""
Basic test for convenience function converting RA/DEC to SEP/PA
"""
ra = np.array([-1., 0., -1., 1.])
dec = np.array([0., -1., -1., 1.])
sep, pa = system.radec2seppa(ra, dec)
assert sep.all() == np.array([1., 1., np.sqrt(2.), np.sqrt(2.)]).all()
assert pa.all() == np.array([270., 180., 225., 45.]).all()
def test_radec2seppa():
ras = np.array([-1, -1, 1, 1])
decs = np.array([-1, 1, -1, 1])
pas_expected = np.array([225., 315., 135., 45.])
pas_expected_180mod = np.array([225., 315., 495., 405.])
seps_expected = np.ones(4)*np.sqrt(2)
sep_nomod, pa_nomod = system.radec2seppa(ras, decs)
sep_180mod, pa_180mod = system.radec2seppa(ras, decs, mod180=True)
assert sep_nomod == pytest.approx(seps_expected, abs=1e-3)
assert sep_180mod == pytest.approx(seps_expected, abs=1e-3)
assert pa_nomod == pytest.approx(pas_expected, abs=1e-3)
assert pa_180mod == pytest.approx(pas_expected_180mod, abs=1e-3)
if min_epoch is None:
# Don't need to rotate and scale if no astrometric measurments for this body. Brute force rejection sampling
continue
period_prescale = np.sqrt(
4*np.pi**2*(sma*u.AU)**3/(consts.G*(mtot*u.Msun))
)
period_prescale = period_prescale.to(u.day).value
meananno = self.epochs[min_epoch]/period_prescale - tau
# compute sep/PA of generated orbits
ra, dec, vc = orbitize.kepler.calc_orbit(
self.epochs[min_epoch], sma, ecc, inc, argp, lan, tau, plx, mtot,
mass_for_Kamp=m1
)
sep, pa = orbitize.system.radec2seppa(ra, dec) # sep[mas], PA[deg]
# generate Gaussian offsets from observational uncertainties
sep_offset = np.random.normal(
0, self.sep_err[min_epoch], size=num_samples
)
pa_offset = np.random.normal(
0, self.pa_err[min_epoch], size=num_samples
)
# calculate correction factors
sma_corr = (sep_offset + self.sep_observed[min_epoch])/sep
lan_corr = (pa_offset + self.pa_observed[min_epoch] - pa)
# perform scale-and-rotate
sma *= sma_corr # [AU]
lan += np.radians(lan_corr) # [rad]