How to use the orbitize.kepler function in orbitize

# 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

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

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
    t = table.Table([epochs, np.ones(epochs.shape, dtype=int), ra_model_b, 0.00001 * np.ones(epochs.shape, dtype=int), dec_model_b, 0.00001 * np.ones(epochs.shape, dtype=int)], 
                     names=["epoch", "object" ,"raoff", "raoff_err","decoff","decoff_err"])
    # add c

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

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]
rv_sar_epoch = s.system.data_table[np.where(
    s.system.data_table['quant_type'] == 'rv')][s.epoch_rv_idx[1]]
pdb.set_trace()
# assert sep_sar == pytest.approx(sar_epoch['quant1'], abs=sar_epoch['quant1_err'])  # issue here
#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)

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

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

def test_c_ecc_anom_solver():
    """
    Test the C implementations in orbitize.kepler._calc_ecc_anom() in the iterative and analytical solver regimes by comparing the mean anomaly computed from
    _calc_ecc_anom() output vs the input mean anomaly
    """
    if kepler.cext:
        test_iterative_ecc_anom_solver(use_c = True)
        test_analytical_ecc_anom_solver(use_c = True)

)

                # Calculate ra/dec offsets for all epochs of this orbit
                if rv_time_series:
                    raoff0, deoff0, vzoff0 = kepler.calc_orbit(
                        epochs_seppa[i, :], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind],
                        tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch,
                        mass_for_Kamp=m0[orb_ind]
                    )

                    raoff[i, :] = raoff0
                    deoff[i, :] = deoff0
                    vz_star[i, :] = vzoff0*-(m1[orb_ind]/m0[orb_ind]) + gamma[orb_ind]

                else:
                    raoff0, deoff0, _ = kepler.calc_orbit(
                        epochs_seppa[i, :], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind],
                        tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch,
                    )

                    raoff[i, :] = raoff0
                    deoff[i, :] = deoff0

                yr_epochs = Time(epochs_seppa[i, :], format='mjd').decimalyear
                plot_epochs = np.where(yr_epochs <= sep_pa_end_year)[0]
                yr_epochs = yr_epochs[plot_epochs]

                seps, pas = orbitize.system.radec2seppa(raoff[i, :], deoff[i, :], mod180=mod180)

                plt.sca(ax1)
                plt.plot(yr_epochs, seps, color=sep_pa_color)

vz_star = np.zeros((num_orbits_to_plot, num_epochs_to_plot))
            epochs = np.zeros((num_orbits_to_plot, num_epochs_to_plot))

            # Loop through each orbit to plot and calcualte ra/dec offsets for all points in orbit
            # Need this loops since epochs[] vary for each orbit, unless we want to just plot the same time period for all orbits
            for i in np.arange(num_orbits_to_plot):
                orb_ind = choose[i]
                # Compute period (from Kepler's third law)
                period = np.sqrt(4*np.pi**2.0*(sma*u.AU)**3/(consts.G*(mtot*u.Msun)))
                period = period.to(u.day).value
                # Create an epochs array to plot num_epochs_to_plot points over one orbital period
                epochs[i, :] = np.linspace(start_mjd, float(
                    start_mjd+period[orb_ind]), num_epochs_to_plot)

                # Calculate ra/dec offsets for all epochs of this orbit
                raoff0, deoff0, _ = kepler.calc_orbit(
                    epochs[i, :], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind],
                    tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch,
                )

                raoff[i, :] = raoff0
                deoff[i, :] = deoff0

            # Create a linearly increasing colormap for our range of epochs
            if cbar_param != 'epochs':
                cbar_param_arr = self.post[:, index]
                norm = mpl.colors.Normalize(vmin=np.min(cbar_param_arr),
                                            vmax=np.max(cbar_param_arr))
                norm_yr = mpl.colors.Normalize(vmin=np.min(
                    cbar_param_arr), vmax=np.max(cbar_param_arr))

            elif cbar_param == 'epochs':