How to use the orbitize.DATADIR function in orbitize

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():

# 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
    for eps, ra, dec in zip(epochs, ra_model_c, dec_model_c):
        t.add_row([eps, 2, ra, 0.000001, dec, 0.000001])

    filename = os.path.join(orbitize.DATADIR, "multiplanet_fake_2planettest.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 = system.System(2, astrom_dat, m0, plx, tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True)

    # fix most of the orbital parameters to make the dimensionality a bit smaller
    sys.sys_priors[sys.param_idx['ecc1']] = b_params[1]
    sys.sys_priors[sys.param_idx['inc1']] = b_params[2]
    sys.sys_priors[sys.param_idx['aop1']] = b_params[3]
    sys.sys_priors[sys.param_idx['pan1']] = b_params[4]

    sys.sys_priors[sys.param_idx['ecc2']] = c_params[1]
    sys.sys_priors[sys.param_idx['inc2']] = c_params[2]
    sys.sys_priors[sys.param_idx['aop2']] = c_params[3]
    sys.sys_priors[sys.param_idx['pan2']] = c_params[4]

def test_write_orbitize_input_2():
    """
    Test the write_orbitize_input and the read_orbitize_input functions

    This test exists with the fail_if_not_removed decorator as a reminder to remove in v2.0
    """
    input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
    test_table = read_file(input_file)
    output_file = os.path.join(orbitize.DATADIR, 'temp_test_orbitize_input.csv')
    # If temp output file already exists, delete it
    if os.path.isfile(output_file):
        os.remove(output_file)
    try:  # Catch these tests so that we remove temporary file
        # Test that we were able to write the table
        write_orbitize_input(test_table,output_file)
        assert os.path.isfile(output_file)
        # Test that we can read the table and check if it's correct
        test_table_2 = read_orbitize_input(output_file)
        _compare_table(test_table_2)
    finally:
        # Remove temporary file
        os.remove(output_file)

def test_ensemble_mcmc_runs(num_threads=1):
    """
    Tests the EnsembleMCMC sampler by making sure it even runs
    """

    # use the test_csv dir
    input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')

    myDriver = Driver(input_file, 'MCMC', 1, 1, 0.01,
                      system_kwargs={'fit_secondary_mass': True,
                                     'tau_ref_epoch': 0},
                      mcmc_kwargs={'num_temps': 1, 'num_threads': num_threads, 'num_walkers': 100}
                      )

    # run it a little (tests 0 burn-in steps)
    myDriver.sampler.run_sampler(100)

    # run it a little more
    myDriver.sampler.run_sampler(1000, burn_steps=1)

    # run it a little more (tests adding to results object)
    myDriver.sampler.run_sampler(1000, burn_steps=1)

def test_custom_likelihood():
    """
    Tests the inclusion of a custom likelihood function in the code
    """
    # use the test_csv dir
    testdir = orbitize.DATADIR
    input_file = os.path.join(testdir, 'GJ504.csv')
    data_table = read_input.read_file(input_file)
    # Manually set 'object' column of data table
    data_table['object'] = 1

    # construct the system
    orbit = system.System(1, data_table, 1, 0.01)

    # construct custom likelihood function
    def my_likelihood(params):
        return -5

    # construct sampler
    n_walkers = 100
    mcmc1 = sampler.MCMC(orbit, 0, n_walkers, num_threads=1)
    mcmc2 = sampler.MCMC(orbit, 0, n_walkers, num_threads=1, custom_lnlike=my_likelihood)

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_compute_model():
    """
    Test basic functionality of ``System.compute_model()``
    """
    input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
    data_table = read_input.read_file(input_file)
    data_table['object'] = 1
    testSystem_parsing = system.System(
        1, data_table, 10., 10.
    )

    params_arr = np.array([[1., 0.5], [0., 0.], [0., 0.], [0., 0.], [
                          0., 0.], [245000., 245000.], [10, 10], [10, 10]])
    model, jitter = testSystem_parsing.compute_model(params_arr)
    assert model.shape == (4, 2, 2)

    params_arr = np.array([1., 0., 0., 0., 0., 245000., 10, 10])
    model, jitter = testSystem_parsing.compute_model(params_arr)
    assert model.shape == (4, 2)

def test_write_orbitize_input_2():
    """
    Test the write_orbitize_input and the read_orbitize_input functions

    This test exists with the fail_if_not_removed decorator as a reminder to remove in v2.0
    """
    input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
    test_table = read_file(input_file)
    output_file = os.path.join(orbitize.DATADIR, 'temp_test_orbitize_input.csv')
    # If temp output file already exists, delete it
    if os.path.isfile(output_file):
        os.remove(output_file)
    try:  # Catch these tests so that we remove temporary file
        # Test that we were able to write the table
        write_orbitize_input(test_table,output_file)
        assert os.path.isfile(output_file)
        # Test that we can read the table and check if it's correct
        test_table_2 = read_orbitize_input(output_file)
        _compare_table(test_table_2)
    finally:
        # Remove temporary file
        os.remove(output_file)

def test_create_driver_from_table():
    """
    Test creation of Driver object from Table as input
    """
    input_file = os.path.join(orbitize.DATADIR, 'test_val.csv')
    input_table = read_file(input_file)
    myDriver = driver.Driver(input_table, # astropy.table Table of input
                             'MCMC', # name of algorithm 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) # parallax error [mas]
    _compare_table(myDriver.system.data_table)