<commit_before>import os
import sys
CUR_PATH = os.path.abspath(os.path.dirname(__file__))
sys.path.append(os.path.join(CUR_PATH, "../../"))
import pytest
import tempfile
import cPickle as pickle
import numpy as np
import os

'''
------------------------------------------------------------------------
Setup
------------------------------------------------------------------------
get_baseline = Flag to run baseline or tax experiments (bool)
calibrate_model = Flag to run calibration of chi values or not (bool)
------------------------------------------------------------------------
'''

TEST_OUTPUT = "./TESTOUT"

newfiles = [os.path.join(TEST_OUTPUT, "SSinit/ss_init_tpi_vars.pkl"),
            os.path.join(TEST_OUTPUT, "SSinit/ss_init_vars.pkl"),
            os.path.join(TEST_OUTPUT, "Saved_moments/SS_init_solutions.pkl"),
            os.path.join(TEST_OUTPUT, "Saved_moments/labor_data_moments.pkl"),
            os.path.join(TEST_OUTPUT, "Saved_moments/wealth_data_moments.pkl"),
            os.path.join(TEST_OUTPUT, "TPIinit/TPIinit_vars.pkl")]

oldfiles = [os.path.join(CUR_PATH, "../../../OUTPUT/SSinit/ss_init_tpi_vars.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/SSinit/ss_init_vars.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/Saved_moments/SS_init_solutions.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/Saved_moments/labor_data_moments.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/Saved_moments/wealth_data_moments.pkl"),
            os.path.join(CUR_PATH, "../../../OUTPUT/TPIinit/TPIinit_vars.pkl")]


@pytest.mark.full_run
def test_sstpi():
    import tempfile
    import pickle
    import numpy as np
    import numpy as np
    import cPickle as pickle
    import os

    import ogusa
    ogusa.parameters.DATASET = 'REAL'

    from ogusa.utils import comp_array
    from ogusa.utils import comp_scalar
    from ogusa.utils import dict_compare
    from ogusa.utils import pickle_file_compare

    import ogusa.SS
    import ogusa.TPI
    from ogusa import parameters, wealth, labor, demographics, income, SS, TPI

    globals().update(ogusa.parameters.get_parameters())

    # Generate Wealth data moments
    output_dir = TEST_OUTPUT
    input_dir = "./OUTPUT"
    wealth.get_wealth_data(lambdas, J, flag_graphs, output_dir)

    # Generate labor data moments
    labor.labor_data_moments(flag_graphs, output_dir=output_dir)

    get_baseline = True
    calibrate_model = False

    # List of parameter names that will not be changing (unless we decide to
    # change them for a tax experiment)
    param_names = ['S', 'J', 'T', 'lambdas', 'starting_age', 'ending_age',
                   'beta', 'sigma', 'alpha', 'nu', 'Z', 'delta', 'E',
                   'ltilde', 'g_y', 'maxiter', 'mindist_SS', 'mindist_TPI',
                   'b_ellipse', 'k_ellipse', 'upsilon',
                   'a_tax_income', 'chi_b_guess', 'chi_n_guess',
                   'b_tax_income', 'c_tax_income', 'd_tax_income',
                   'tau_payroll', 'tau_bq', 'calibrate_model',
                   'retire', 'mean_income_data', 'g_n_vector',
                   'h_wealth', 'p_wealth', 'm_wealth', 'get_baseline',
                   'omega', 'g_n_ss', 'omega_SS', 'surv_rate', 'e', 'rho']

    '''
    ------------------------------------------------------------------------
        Run SS with minimization to fit chi_b and chi_n
    ------------------------------------------------------------------------
    '''

    # This is the simulation before getting the replacement rate values
    sim_params = {}
    for key in param_names:
        try:
            sim_params[key] = locals()[key]
        except KeyError:
            sim_params[key] = globals()[key]<commit_msg>edits test of wealth.py and labor.py as functions changed<commit_after>import os
import sys
CUR_PATH = os.path.abspath(os.path.dirname(__file__))
sys.path.append(os.path.join(CUR_PATH, "../../"))
import pytest
import tempfile
import cPickle as pickle
import numpy as np
import os

'''
------------------------------------------------------------------------
Setup
------------------------------------------------------------------------
get_baseline = Flag to run baseline or tax experiments (bool)
calibrate_model = Flag to run calibration of chi values or not (bool)
------------------------------------------------------------------------
'''

TEST_OUTPUT = "./TESTOUT"

newfiles = [os.path.join(TEST_OUTPUT, "SSinit/ss_init_tpi_vars.pkl"),
            os.path.join(TEST_OUTPUT, "SSinit/ss_init_vars.pkl"),
            os.path.join(TEST_OUTPUT, "Saved_moments/SS_init_solutions.pkl"),
            os.path.join(TEST_OUTPUT, "Saved_moments/labor_data_moments.pkl"),
            os.path.join(TEST_OUTPUT, "Saved_moments/wealth_data_moments.pkl"),
            os.path.join(TEST_OUTPUT, "TPIinit/TPIinit_vars.pkl")]

oldfiles = [os.path.join(CUR_PATH, "../../../OUTPUT/SSinit/ss_init_tpi_vars.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/SSinit/ss_init_vars.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/Saved_moments/SS_init_solutions.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/Saved_moments/labor_data_moments.pkl"),
            os.path.join(
                CUR_PATH, "../../../OUTPUT/Saved_moments/wealth_data_moments.pkl"),
            os.path.join(CUR_PATH, "../../../OUTPUT/TPIinit/TPIinit_vars.pkl")]


@pytest.mark.full_run
def test_sstpi():
    import tempfile
    import pickle
    import numpy as np
    import numpy as np
    import cPickle as pickle
    import os

    import ogusa
    ogusa.parameters.DATASET = 'REAL'

    from ogusa.utils import comp_array
    from ogusa.utils import comp_scalar
    from ogusa.utils import dict_compare
    from ogusa.utils import pickle_file_compare

    import ogusa.SS
    import ogusa.TPI
    from ogusa import parameters, wealth, labor, demographics, income, SS, TPI

    globals().update(ogusa.parameters.get_parameters())

    # Generate Wealth data moments
    output_dir = TEST_OUTPUT
    input_dir = "./OUTPUT"
    wealth.get_wealth_data()

    # Generate labor data moments

    get_baseline = True
    calibrate_model = False

    # List of parameter names that will not be changing (unless we decide to
    # change them for a tax experiment)
    param_names = ['S', 'J', 'T', 'lambdas', 'starting_age', 'ending_age',
                   'beta', 'sigma', 'alpha', 'nu', 'Z', 'delta', 'E',
                   'ltilde', 'g_y', 'maxiter', 'mindist_SS', 'mindist_TPI',
                   'b_ellipse', 'k_ellipse', 'upsilon',
                   'a_tax_income', 'chi_b_guess', 'chi_n_guess',
                   'b_tax_income', 'c_tax_income', 'd_tax_income',
                   'tau_payroll', 'tau_bq', 'calibrate_model',
                   'retire', 'mean_income_data', 'g_n_vector',
                   'h_wealth', 'p_wealth', 'm_wealth', 'get_baseline',
                   'omega', 'g_n_ss', 'omega_SS', 'surv_rate', 'e', 'rho']

    '''
    ------------------------------------------------------------------------
        Run SS with minimization to fit chi_b and chi_n
    ------------------------------------------------------------------------
    '''

    # This is the simulation before getting the replacement rate values
    sim_params = {}
    for key in param_names:
        try:
            sim_params[key] = locals()[key]
        except KeyError:
            sim_params[key] = globals()[key]
{'fraction_non_alphanumeric': 0.2120010573618821, 'fraction_numerical': 0.0, 'mean_word_length': 1.8393996247654785, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5799399', 'n_tokens_mistral': 2350, 'n_tokens_neox': 2257, 'n_words': 494}