How to use orbitize - 10 common examples

# generate planet c orbital parameters
    # at 2 au, and starts on the opposite side of the star relative to b
    c_params = [2, 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_c*365.25, 20) + tau_ref_epoch # the full period of c, 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 b due to c
    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

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]
    inc = samples[:, 2]
    argp = samples[:, 3]
    lan = samples[:, 4]
    tau = samples[:, 5]
    plx = samples[:, 6]
    mtot = samples[:, 7]

    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
    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'])

ecc = samples[:, 1]
    inc = samples[:, 2]
    argp = samples[:, 3]
    lan = samples[:, 4]
    tau = samples[:, 5]
    plx = samples[:, 6]
    if self.system.fit_secondary_mass:
        gamma = samples[:, 7]
        sigma = samples[:, 8]
        m0 = samples[:, -1]
        m1 = samples[:, -2]
        mtot = m0 + m1
    else:
        mtot = samples[:, 7]

    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
    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_orbit_e03_array():
    """
    Test orbitize.kepler.calc_orbit() with a standard orbit with ecc = 0.3 and an array of keplerian input
    """
    # sma, ecc, inc, argp, lan, tau, plx, mtot
    sma = np.array([10,10,10])
    ecc = np.array([0.3,0.3,0.3])
    inc = np.array([3,3,3])
    argp = np.array([0.5,0.5,0.5])
    lan = np.array([1.5,1.5,1.5])
    tau = np.array([0.3,0.3,0.3])
    plx = np.array([50,50,50])
    mtot = np.array([1.5,1.5,1.5])
    epochs = np.array([1000, 1101.4])
    raoffs, deoffs, vzs = kepler.calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot)

    true_raoff = np.array([[ 152.86786, 152.86786, 152.86786],
                           [ 180.39408, 180.39408, 180.39408]])
    true_deoff = np.array([[-462.91038,-462.91038,-462.91038],
                           [-442.0127, -442.0127, -442.0127]])
    true_vz    = np.array([[.86448656, .86448656, .86448656],
                           [.97591289, .97591289, .97591289]])

    for ii in range(0,3):
        for meas, truth in zip(raoffs[:, ii], true_raoff[:,ii]):
            assert truth == pytest.approx(meas, abs=threshold)
        for meas, truth in zip(deoffs[:, ii], true_deoff[:, ii]):
            assert truth == pytest.approx(meas, abs=threshold)
        for meas, truth in zip(vzs[:, ii], true_vz[:, ii]):
            assert truth == pytest.approx(meas, abs=1e-8)

def test_orbit_e99():
    """
    Test a highly eccentric orbit (ecc=0.99). Again validate against James Graham's orbit code
    """
    # sma, ecc, inc, argp, lan, tau, plx, mtot
    orbital_params = np.array([10, 0.99, 3, 0.5, 1.5, 0.3, 50, 1.5])
    epochs = np.array([1000, 1101.4])
    raoffs, deoffs, vzs = kepler.calc_orbit(epochs, orbital_params[0], orbital_params[1], orbital_params[2], orbital_params[3],
                                        orbital_params[4], orbital_params[5], orbital_params[6], orbital_params[7])

    true_raoff = [-589.45575, -571.48432]
    true_deoff = [-447.32217, -437.68456]
    true_vz = [.39208876,  .42041953]

    for meas, truth in zip(raoffs, true_raoff):
        assert truth == pytest.approx(meas, abs=threshold)
    for meas, truth in zip(deoffs, true_deoff):
        assert truth == pytest.approx(meas, abs=threshold)
    for meas, truth in zip(vzs, true_vz):
        assert truth == pytest.approx(meas, abs=1e-8)

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]
    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])

rv_sar = np.median(v_star[s.epoch_rv_idx[1]])


# test to make sure sep and pa scaled to scale-and-rotate epoch
sar_epoch = s.system.data_table[np.where(
    s.system.data_table['quant_type'] == 'seppa')][s.epoch_idx]

def test_systeminit():
    """
    Test that initializing a ``System`` class produces a list of ``Prior``
    objects of the correct length when:
        - parallax and total mass are fixed
        - parallax and total mass errors are given
        - parallax is fixed, total mass error is given
        - parallax error is given, total mass error is fixed

    Test that the different types of data are parsed correctly
    when initializing a ``System`` object.
    """
    testdir = orbitize.DATADIR
    input_file = os.path.join(testdir, 'test_val.csv')
    data_table = read_input.read_file(input_file)

    # Manually set 'object' column of data table
    data_table['object'] = 1
    data_table['object'][1] = 2

    plx_mass_errs2lens = {
        (0., 0.): 14,
        (1., 1.): 14,
        (0., 1.): 14,
        (1., 0.): 14
    }

    for plx_e, mass_e in plx_mass_errs2lens.keys():

        testSystem_priors = system.System(
            2, data_table, 10., 10., plx_err=plx_e, mass_err=mass_e

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]

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'])