How to use the orbitize.kepler.calc_orbit function in orbitize

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

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]

if self.system.fit_secondary_mass:
            m0 = samples[-1,:]
            m1 = samples[-2,:]
            mtot = m0 + m1
        else:
            mtot = samples[-1,:]
            m1 = None

        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[self.epoch_idx]/period_prescale - tau

        # compute sep/PA of generated orbits
        ra, dec, vc = orbitize.kepler.calc_orbit(
            self.epochs[self.epoch_idx], 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[self.epoch_idx], size=num_samples
        )
        pa_offset =  np.random.normal(
            0, self.pa_err[self.epoch_idx], size=num_samples
        )

        # calculate correction factors
        sma_corr = (sep_offset + self.sep_observed[self.epoch_idx])/sep
        lan_corr = (pa_offset + self.pa_observed[self.epoch_idx] - pa)