<commit_before>'''
------------------------------------------------------------------------
Last updated 4/8/2015

This file sets parameters for the model run.

This py-file calls the following other file(s):
            income.py
            demographics.py
------------------------------------------------------------------------
'''

'''
------------------------------------------------------------------------
Import Packages
------------------------------------------------------------------------
'''
import os
import json
import numpy as np
import scipy.ndimage.filters as filter
from demographics import get_pop_objs
from demographics_old import get_omega
from income import get_e
import pickle
import txfunc
import elliptical_u_est
import matplotlib.pyplot as plt


'''
------------------------------------------------------------------------
Set paths, define user modifiable parameters
------------------------------------------------------------------------
'''
DATASET = 'REAL'
PARAMS_FILE = os.path.join(os.path.dirname(__file__), 'default_full_parameters.json')
PARAMS_FILE_METADATA_NAME = 'parameters_metadata.json'
PARAMS_FILE_METADATA_PATH = os.path.join(os.path.dirname(__file__), PARAMS_FILE_METADATA_NAME)
TAX_ESTIMATE_PATH = os.environ.get("TAX_ESTIMATE_PATH", ".")
USER_MODIFIABLE_PARAMS = ['g_y_annual', 'frisch']


def read_parameter_metadata():
    '''
    --------------------------------------------------------------------
    This function reads in parameter metadata
    --------------------------------------------------------------------

    INPUTS: None

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    /PARAMS_FILE_METADATA_PATH/ = json file with metadata

    OBJECTS CREATED WITHIN FUNCTION:
    params_dict = dictionary of metadata

    RETURNS: params_dict
    --------------------------------------------------------------------
    '''
    if os.path.exists(PARAMS_FILE_METADATA_PATH):
        with open(PARAMS_FILE_METADATA_PATH) as pfile:
            params_dict = json.load(pfile)
    else:
        from pkg_resources import resource_stream, Requirement
        path_in_egg = os.path.join('ogusa', PARAMS_FILE_METADATA_NAME)
        buf = resource_stream(Requirement.parse('ogusa'), path_in_egg)
        as_bytes = buf.read()
        as_string = as_bytes.decode("utf-8")
        params_dict = json.loads(as_string)

    return params_dict


def read_tax_func_estimate(pickle_path, pickle_file):
    '''
    --------------------------------------------------------------------
    This function reads in tax function parameters
    --------------------------------------------------------------------

    INPUTS:
    pickle_path = string, path to pickle with tax function parameter estimates
    pickle_file = string, name of pickle file with tax function parmaeter estimates

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    /picklepath/ = pickle file with dictionary of tax function estimated parameters

    OBJECTS CREATED WITHIN FUNCTION:
    dict_params = dictionary, contains numpy arrays of tax function estimates

    RETURNS: dict_params
    --------------------------------------------------------------------
    '''
    if os.path.exists(pickle_path):
        with open(pickle_path) as pfile:
            dict_params = pickle.load(pfile)
    else:
        from pkg_resources import resource_stream, Requirement
        path_in_egg = pickle_file
        buf = resource_stream(Requirement.parse('ogusa'), path_in_egg)
        as_bytes = buf.read()
        as_string = as_bytes.decode("utf-8")
        dict_params = pickle.loads(as_string)

    return dict_params


def get_parameters_from_file():
    '''
    --------------------------------------------------------------------
    This function loads the json file with model parameters
    --------------------------------------------------------------------

    INPUTS: None

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    /PARAMS_FILE/ = json file with model parameters

    OBJECTS CREATED WITHIN FUNCTION:

    RETURNS: j
    --------------------------------------------------------------------
    '''
    with open(PARAMS_FILE,'r') as f:
        j = json.load(f)
        for key in j:
            if isinstance(j[key], list):
                j[key] = np.array(j[key])
        return j


def get_parameters(baseline=False, guid='', user_modifiable=False, metadata=False):
    '''
    --------------------------------------------------------------------
    This function returns the model parameters.
    --------------------------------------------------------------------

    INPUTS:
    baseline        = boolean, =True if run is of baseline policy
    guid            = string, id for model run
    user_modifiable = boolean, =True if allow user modifiable parameters
    metadata        = boolean, =True if use metadata file for parameter
                       values (rather than what is entered in parameters below)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    get_full_parameters()
    get_reduced_parameters()

    OBJECTS CREATED WITHIN FUNCTION:

    RETURNS: dictionary with model parameters
    --------------------------------------------------------------------
    '''
    if DATASET == 'REAL':
        return get_full_parameters(baseline=baseline, guid=guid,
                                   user_modifiable=user_modifiable,
                                   metadata=metadata)

    elif DATASET == 'SMALL':
        return get_reduced_parameters(baseline=baseline, guid=guid,
                                      user_modifiable=user_modifiable,
                                      metadata=metadata)
    else:
        raise ValueError("Unknown value {0}".format(DATASET))


'''
------------------------------------------------------------------------
Parameters
------------------------------------------------------------------------
Model Parameters:
------------------------------------------------------------------------
S            = integer, number of economically active periods an individual lives
J            = integer, number of different ability groups
T            = integer, number of time periods until steady state is reached
BW           = integer, number of time periods in the budget window
lambdas      = [J,] vector, percentiles for ability groups
imm_rates    = [J,T+S] array, immigration rates by age and year
starting_age = integer, age agents enter population
ending age   = integer, maximum age agents can live until
E            = integer, age agents become economically active
beta_annual  = scalar, discount factor as an annual rate
beta         = scalar, discount factor for model period
sigma        = scalar, coefficient of relative risk aversion
alpha        = scalar, capital share of income
Z            = scalar, total factor productivity parameter in firms' production
               function
delta_annual = scalar, depreciation rate as an annual rate
delta        = scalar, depreciation rate for model period
ltilde       = scalar, measure of time each individual is endowed with each
               period
g_y_annual   = scalar, annual growth rate of technology
g_y          = scalar, growth rate of technology for a model period
frisch       = scalar, Frisch elasticity that is used to fit ellipitcal utility
               to constant Frisch elasticity function
b_ellipse    = scalar, value of b for elliptical fit of utility function
k_ellipse    = scalar, value of k for elliptical fit of utility function
upsilon      = scalar, value of omega for elliptical fit of utility function
------------------------------------------------------------------------
Tax Parameters:
------------------------------------------------------------------------
mean_income_data = scalar, mean income from IRS data file used to calibrate income tax
etr_params       = [S,BW,#tax params] array, parameters for effective tax rate function
mtrx_params      = [S,BW,#tax params] array, parameters for marginal tax rate on
                    labor income function
mtry_params      = [S,BW,#tax params] array, parameters for marginal tax rate on
                    capital income function
h_wealth         = scalar, wealth tax parameter h (scalar)
m_wealth         = scalar, wealth tax parameter m (scalar)
p_wealth         = scalar, wealth tax parameter p (scalar)
tau_bq           = [J,] vector, bequest tax
tau_payroll      = scalar, payroll tax rate
retire           = integer, age at which individuals eligible for retirement benefits
------------------------------------------------------------------------
Simulation Parameters:
------------------------------------------------------------------------
MINIMIZER_TOL = scalar, tolerance level for the minimizer in the calibration of chi parameters
MINIMIZER_OPTIONS = dictionary, dictionary for options to put into the minimizer, usually
                    to set a max iteration
PLOT_TPI     = boolean, =Ture if plot the path of K as TPI iterates (for debugging purposes)
maxiter      = integer, maximum number of iterations that SS and TPI solution methods will undergo
mindist_SS   = scalar, tolerance for SS solution
mindist_TPI  = scalar, tolerance for TPI solution
nu           = scalar, contraction parameter in SS and TPI iteration process
               representing the weight on the new distribution
flag_graphs  = boolean, =True if produce graphs in demographic, income,
               wealth, and labor files (True=graph)
chi_b_guess  = [J,] vector, initial guess of \chi^{b}_{j} parameters
               (if no calibration occurs, these are the values that will be used for \chi^{b}_{j})
chi_n_guess  = [S,] vector, initial guess of \chi^{n}_{s} parameters
               (if no calibration occurs, these are the values that will be used for \chi^{n}_{s})
------------------------------------------------------------------------
Demographics and Ability variables:
------------------------------------------------------------------------
omega        = [T+S,S] array, time path of stationary distribution of economically active population by age
g_n_ss       = scalar, steady state population growth rate
omega_SS     = [S,] vector, stationary steady state population distribution
surv_rate    = [S,] vector, survival rates by age
rho          = [S,] vector, mortality rates by age
g_n_vector   = [T+S,] vector, growth rate in economically active pop for each period in transition path
e            = [S,J] array, normalized effective labor units by age and ability type
------------------------------------------------------------------------
'''


def get_reduced_parameters(baseline, guid, user_modifiable, metadata):
    '''
    --------------------------------------------------------------------
    This function sets the parameters for the reduced model, which is
    simplified to run more quickly for testing.
    --------------------------------------------------------------------

    INPUTS:
    baseline        = boolean, =True if baseline tax policy, =False if reform
    guid            = string, id for reform run
    user_modifiable = boolean, =True if allow user modifiable parameters
    metadata        = boolean, =True if use metadata file for parameter
                       values (rather than what is entered in parameters below)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    read_tax_func_estimate()
    elliptical_u_est.estimation()
    read_parameter_metadata()

    OBJECTS CREATED WITHIN FUNCTION:
    See parameters defined above
    allvars = dictionary, dictionary with all parameters defined in this function

    RETURNS: allvars

    OUTPUT: None
    --------------------------------------------------------------------
    '''
    # Model Parameters
    starting_age = 40
    ending_age = 50
    S = int(ending_age-starting_age)
    J = int(1)
    T = int(2 * S)
    BW = int(10)
    lambdas = np.array([1.])
    E = int(starting_age * (S / float(ending_age - starting_age)))
    beta_annual = .96 # Carroll (JME, 2009)
    beta = beta_annual ** (float(ending_age - starting_age) / S)
    sigma = 1.5 # value from Attanasio, Banks, Meghir and Weber (JEBS, 1999)
    alpha = .35 # many use 0.33, but many find that capitals share is increasing (e.g. Elsby, Hobijn, and Sahin (BPEA, 2013))
    Z = 1.0
    delta_annual = .05 # approximately the value from Kehoe calibration exercise: http://www.econ.umn.edu/~tkehoe/classes/calibration-04.pdf
    delta = 1 - ((1 - delta_annual) ** (float(ending_age - starting_age) / S))
    ltilde = 1.0
    g_y_annual = 0.03
    g_y = (1 + g_y_annual)**(float(ending_age - starting_age) / S) - 1
    #   Ellipse parameters
    frisch = 0.4 # Frisch elasticity consistent with Altonji (JPE, 1996) and Peterman (Econ Inquiry, 2016)
    b_ellipse, upsilon = elliptical_u_est.estimation(frisch,ltilde)
    k_ellipse = 0 # this parameter is just a level shifter in utlitiy - irrelevant for analysis

    # Tax parameters:
    #   Income Tax Parameters
    ####  will call tax function estimation function here...
    ### do output such that each parameters is in a separate SxBW array

    if baseline:
        baseline_pckl = "TxFuncEst_baseline{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     baseline_pckl)
        print 'using baseline1 tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, baseline_pckl)
    else:
        policy_pckl = "TxFuncEst_policy{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     policy_pckl)
        print 'using policy1 tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, policy_pckl)


    mean_income_data = dict_params['tfunc_avginc'][0]

    etr_params = dict_params['tfunc_etr_params_S'][:S,:BW,:]
    mtrx_params = dict_params['tfunc_mtrx_params_S'][:S,:BW,:]
    mtry_params = dict_params['tfunc_mtry_params_S'][:S,:BW,:]

    # To zero out income taxes, uncomment the following 3 lines:
    # etr_params[:,:,6:] = 0.0
    # mtrx_params[:,:,6:] = 0.0
    # mtry_params[:,:,6:] = 0.0

    #   Wealth tax params
    #       These are non-calibrated values, h and m just need
    #       need to be nonzero to avoid errors. When p_wealth
    #       is zero, there is no wealth tax.
    h_wealth = 0.1
    m_wealth = 1.0
    p_wealth = 0.0
    #   Bequest and Payroll Taxes
    tau_bq = np.zeros(J)
    tau_payroll = 0.15
    retire = np.round(9.0 * S / 16.0) - 1

    # Simulation Parameters
    MINIMIZER_TOL = 1e-3
    MINIMIZER_OPTIONS = {'maxiter': 1}
    PLOT_TPI = False
    maxiter = 10
    mindist_SS = 1e-3
    mindist_TPI = 1e-3 #1e-6
    nu = .4
    flag_graphs = False
    #   Calibration parameters
    # These guesses are close to the calibrated values
    chi_b_guess = np.array([1, 100000])
    chi_n_guess = np.array([5, 6, 7, 8, 9, 10, 11, 12, 13, 14])

    # Generate Income and Demographic parameters
    omega, g_n_ss, omega_SS, surv_rate, rho, g_n_vector, imm_rates, omega_S_preTP = get_pop_objs(
        E, S, T, 1, 100, 2016, flag_graphs)
    e = np.array([[0.25, 1.25]] * 10)
    allvars = dict(locals())

    if user_modifiable:
        allvars = {k:allvars[k] for k in USER_MODIFIABLE_PARAMS}

    if metadata:
        params_meta = read_parameter_metadata()
        for k,v in allvars.iteritems():
            params_meta[k]["value"] = v
        allvars = params_meta

    return allvars


def get_full_parameters(baseline, guid, user_modifiable, metadata):
    '''
    --------------------------------------------------------------------
    This function sets the parameters for the full model.
    --------------------------------------------------------------------

    INPUTS:
    baseline        = boolean, =True if baseline tax policy, =False if reform
    guid            = string, id for reform run
    user_modifiable = boolean, =True if allow user modifiable parameters
    metadata        = boolean, =True if use metadata file for parameter
                       values (rather than what is entered in parameters below)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    read_tax_func_estimate()
    elliptical_u_est.estimation()
    read_parameter_metadata()

    OBJECTS CREATED WITHIN FUNCTION:
    See parameters defined above
    allvars = dictionary, dictionary with all parameters defined in this function

    RETURNS: allvars

    OUTPUT: None
    --------------------------------------------------------------------
    '''
    # Model Parameters
    S = int(30) #S<30 won't meet necessary tolerances
    J = int(1)
    T = int(10 * S)
    BW = int(10)
    #lambdas = np.array([.25, .25, .2, .1, .1, .09, .01])
    #lambdas = np.array([0.5, 0.5])
    lambdas = np.array([1.,])

    starting_age = 20
    ending_age = 100
    E = int(starting_age * (S / float(ending_age - starting_age)))
    beta_annual = .96 # Carroll (JME, 2009)
    beta = beta_annual ** (float(ending_age - starting_age) / S)
    sigma = 1.5 # value from Attanasio, Banks, <NAME> Weber (JEBS, 1999)
    alpha = .35 # many use 0.33, but many find that capitals share is increasing (e.g. Elsby, Hobijn, and Sahin (BPEA, 2013))
    Z = 1.0
    delta_annual = .05 # approximately the value from Kehoe calibration exercise: http://www.econ.umn.edu/~tkehoe/classes/calibration-04.pdf
    delta = 1 - ((1 - delta_annual) ** (float(ending_age - starting_age) / S))
    ltilde = 1.0
    g_y_annual = 0.03
    g_y = (1 + g_y_annual)**(float(ending_age - starting_age) / S) - 1
    #   Ellipse parameters
    frisch = 0.4 # Frisch elasticity consistent with Altonji (JPE, 1996) and Peterman (Econ Inquiry, 2016)
    b_ellipse, upsilon = elliptical_u_est.estimation(frisch,ltilde)
    k_ellipse = 0 # this parameter is just a level shifter in utlitiy - irrelevant for analysis

    # Tax parameters:
    #   Income Tax Parameters
    #  will call tax function estimation function here...
    # do output such that each parameters is in a separate SxBW array
    # read in estimated parameters
    print 'baselines is:', baseline
    if baseline:
        baseline_pckl = "TxFuncEst_baseline{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     baseline_pckl)
        print 'using baseline tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, baseline_pckl)

    else:
        policy_pckl = "TxFuncEst_policy{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     policy_pckl)
        print 'using policy tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, policy_pckl)


    mean_income_data = dict_params['tfunc_avginc'][0]

    etr_params = dict_params['tfunc_etr_params_S'][:S,:BW,:]
    mtrx_params = dict_params['tfunc_mtrx_params_S'][:S,:BW,:]
    mtry_params = dict_params['tfunc_mtry_params_S'][:S,:BW,:]


    etr_params = np.zeros((S,BW,10))
    mtrx_params = np.zeros((S,BW,10))
    mtry_params = np.zeros((S,BW,10))
    etr_params[:,:,:] = dict_params['tfunc_etr_params_S'][20,0,:]
    mtrx_params[:,:,:] = dict_params['tfunc_mtrx_params_S'][20,0,:]
    mtry_params[:,:,:] = dict_params['tfunc_mtry_params_S'][20,0,:]


    mtrx_params[:,:,1] = 0.
    mtrx_params[:,:,2] = 0.
    mtrx_params[:,:,4] = 0.
    mtrx_params[:,:,8] = 0.
    mtrx_params[:,:,9] = 0.

    mtry_params[:,:,0] = 0.
    mtry_params[:,:,2] = 0.
    mtry_params[:,:,3] = 0.
    mtry_params[:,:,6] = 0.
    mtry_params[:,:,7] = 0.
    
    # # unocmmenting the block below ensures no tax rates are negative
    # etr_params[:,:,7] = 0.
    # mtrx_params[:,:,7] = 0.
    # mtry_params[:,:,7] = 0.
    # etr_params[:,:,9] = 0.
    # mtrx_params[:,:,9] = 0.
    # mtry_params[:,:,9] = 0.

    # print 'tax diffs: ', np.absolute(etr_params-etr_params.mean(axis=(0,1))).max()
    # quit()




    # set etrs and mtrs to constant rates over income/age by uncommenting following code block
    etr_params = np.zeros((S,BW,10))
    # mtrx_params = np.zeros((S,BW,10))
    # mtry_params = np.zeros((S,BW,10))
    etr_params[:,:,7] = dict_params['tfunc_avg_etr']
    # mtrx_params[:,:,7] = dict_params['tfunc_avg_mtrx']
    # mtry_params[:,:,7] = dict_params['tfunc_avg_mtry']
    # etr_params[:,:,9] = dict_params['tfunc_avg_etr']
    # mtrx_params[:,:,9] = dict_params['tfunc_avg_mtrx']
    # mtry_params[:,:,9] = dict_params['tfunc_avg_mtry']
    # etr_params[:,:,5] = 1.0
    # mtrx_params[:,:,5] = 1.0
    # mtry_params[:,:,5] = 1.0

    # etr_params = np.zeros((S,BW,10))
    # mtrx_params = np.zeros((S,BW,10))
    # mtry_params = np.zeros((S,BW,10))
    # for i in xrange(S):
    #     etr_params[:,:,7] = 0.005*i<commit_msg>test mtrx and mtry together<commit_after>'''
------------------------------------------------------------------------
Last updated 4/8/2015

This file sets parameters for the model run.

This py-file calls the following other file(s):
            income.py
            demographics.py
------------------------------------------------------------------------
'''

'''
------------------------------------------------------------------------
Import Packages
------------------------------------------------------------------------
'''
import os
import json
import numpy as np
import scipy.ndimage.filters as filter
from demographics import get_pop_objs
from demographics_old import get_omega
from income import get_e
import pickle
import txfunc
import elliptical_u_est
import matplotlib.pyplot as plt


'''
------------------------------------------------------------------------
Set paths, define user modifiable parameters
------------------------------------------------------------------------
'''
DATASET = 'REAL'
PARAMS_FILE = os.path.join(os.path.dirname(__file__), 'default_full_parameters.json')
PARAMS_FILE_METADATA_NAME = 'parameters_metadata.json'
PARAMS_FILE_METADATA_PATH = os.path.join(os.path.dirname(__file__), PARAMS_FILE_METADATA_NAME)
TAX_ESTIMATE_PATH = os.environ.get("TAX_ESTIMATE_PATH", ".")
USER_MODIFIABLE_PARAMS = ['g_y_annual', 'frisch']


def read_parameter_metadata():
    '''
    --------------------------------------------------------------------
    This function reads in parameter metadata
    --------------------------------------------------------------------

    INPUTS: None

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    /PARAMS_FILE_METADATA_PATH/ = json file with metadata

    OBJECTS CREATED WITHIN FUNCTION:
    params_dict = dictionary of metadata

    RETURNS: params_dict
    --------------------------------------------------------------------
    '''
    if os.path.exists(PARAMS_FILE_METADATA_PATH):
        with open(PARAMS_FILE_METADATA_PATH) as pfile:
            params_dict = json.load(pfile)
    else:
        from pkg_resources import resource_stream, Requirement
        path_in_egg = os.path.join('ogusa', PARAMS_FILE_METADATA_NAME)
        buf = resource_stream(Requirement.parse('ogusa'), path_in_egg)
        as_bytes = buf.read()
        as_string = as_bytes.decode("utf-8")
        params_dict = json.loads(as_string)

    return params_dict


def read_tax_func_estimate(pickle_path, pickle_file):
    '''
    --------------------------------------------------------------------
    This function reads in tax function parameters
    --------------------------------------------------------------------

    INPUTS:
    pickle_path = string, path to pickle with tax function parameter estimates
    pickle_file = string, name of pickle file with tax function parmaeter estimates

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    /picklepath/ = pickle file with dictionary of tax function estimated parameters

    OBJECTS CREATED WITHIN FUNCTION:
    dict_params = dictionary, contains numpy arrays of tax function estimates

    RETURNS: dict_params
    --------------------------------------------------------------------
    '''
    if os.path.exists(pickle_path):
        with open(pickle_path) as pfile:
            dict_params = pickle.load(pfile)
    else:
        from pkg_resources import resource_stream, Requirement
        path_in_egg = pickle_file
        buf = resource_stream(Requirement.parse('ogusa'), path_in_egg)
        as_bytes = buf.read()
        as_string = as_bytes.decode("utf-8")
        dict_params = pickle.loads(as_string)

    return dict_params


def get_parameters_from_file():
    '''
    --------------------------------------------------------------------
    This function loads the json file with model parameters
    --------------------------------------------------------------------

    INPUTS: None

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    /PARAMS_FILE/ = json file with model parameters

    OBJECTS CREATED WITHIN FUNCTION:

    RETURNS: j
    --------------------------------------------------------------------
    '''
    with open(PARAMS_FILE,'r') as f:
        j = json.load(f)
        for key in j:
            if isinstance(j[key], list):
                j[key] = np.array(j[key])
        return j


def get_parameters(baseline=False, guid='', user_modifiable=False, metadata=False):
    '''
    --------------------------------------------------------------------
    This function returns the model parameters.
    --------------------------------------------------------------------

    INPUTS:
    baseline        = boolean, =True if run is of baseline policy
    guid            = string, id for model run
    user_modifiable = boolean, =True if allow user modifiable parameters
    metadata        = boolean, =True if use metadata file for parameter
                       values (rather than what is entered in parameters below)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    get_full_parameters()
    get_reduced_parameters()

    OBJECTS CREATED WITHIN FUNCTION:

    RETURNS: dictionary with model parameters
    --------------------------------------------------------------------
    '''
    if DATASET == 'REAL':
        return get_full_parameters(baseline=baseline, guid=guid,
                                   user_modifiable=user_modifiable,
                                   metadata=metadata)

    elif DATASET == 'SMALL':
        return get_reduced_parameters(baseline=baseline, guid=guid,
                                      user_modifiable=user_modifiable,
                                      metadata=metadata)
    else:
        raise ValueError("Unknown value {0}".format(DATASET))


'''
------------------------------------------------------------------------
Parameters
------------------------------------------------------------------------
Model Parameters:
------------------------------------------------------------------------
S            = integer, number of economically active periods an individual lives
J            = integer, number of different ability groups
T            = integer, number of time periods until steady state is reached
BW           = integer, number of time periods in the budget window
lambdas      = [J,] vector, percentiles for ability groups
imm_rates    = [J,T+S] array, immigration rates by age and year
starting_age = integer, age agents enter population
ending age   = integer, maximum age agents can live until
E            = integer, age agents become economically active
beta_annual  = scalar, discount factor as an annual rate
beta         = scalar, discount factor for model period
sigma        = scalar, coefficient of relative risk aversion
alpha        = scalar, capital share of income
Z            = scalar, total factor productivity parameter in firms' production
               function
delta_annual = scalar, depreciation rate as an annual rate
delta        = scalar, depreciation rate for model period
ltilde       = scalar, measure of time each individual is endowed with each
               period
g_y_annual   = scalar, annual growth rate of technology
g_y          = scalar, growth rate of technology for a model period
frisch       = scalar, Frisch elasticity that is used to fit ellipitcal utility
               to constant Frisch elasticity function
b_ellipse    = scalar, value of b for elliptical fit of utility function
k_ellipse    = scalar, value of k for elliptical fit of utility function
upsilon      = scalar, value of omega for elliptical fit of utility function
------------------------------------------------------------------------
Tax Parameters:
------------------------------------------------------------------------
mean_income_data = scalar, mean income from IRS data file used to calibrate income tax
etr_params       = [S,BW,#tax params] array, parameters for effective tax rate function
mtrx_params      = [S,BW,#tax params] array, parameters for marginal tax rate on
                    labor income function
mtry_params      = [S,BW,#tax params] array, parameters for marginal tax rate on
                    capital income function
h_wealth         = scalar, wealth tax parameter h (scalar)
m_wealth         = scalar, wealth tax parameter m (scalar)
p_wealth         = scalar, wealth tax parameter p (scalar)
tau_bq           = [J,] vector, bequest tax
tau_payroll      = scalar, payroll tax rate
retire           = integer, age at which individuals eligible for retirement benefits
------------------------------------------------------------------------
Simulation Parameters:
------------------------------------------------------------------------
MINIMIZER_TOL = scalar, tolerance level for the minimizer in the calibration of chi parameters
MINIMIZER_OPTIONS = dictionary, dictionary for options to put into the minimizer, usually
                    to set a max iteration
PLOT_TPI     = boolean, =Ture if plot the path of K as TPI iterates (for debugging purposes)
maxiter      = integer, maximum number of iterations that SS and TPI solution methods will undergo
mindist_SS   = scalar, tolerance for SS solution
mindist_TPI  = scalar, tolerance for TPI solution
nu           = scalar, contraction parameter in SS and TPI iteration process
               representing the weight on the new distribution
flag_graphs  = boolean, =True if produce graphs in demographic, income,
               wealth, and labor files (True=graph)
chi_b_guess  = [J,] vector, initial guess of \chi^{b}_{j} parameters
               (if no calibration occurs, these are the values that will be used for \chi^{b}_{j})
chi_n_guess  = [S,] vector, initial guess of \chi^{n}_{s} parameters
               (if no calibration occurs, these are the values that will be used for \chi^{n}_{s})
------------------------------------------------------------------------
Demographics and Ability variables:
------------------------------------------------------------------------
omega        = [T+S,S] array, time path of stationary distribution of economically active population by age
g_n_ss       = scalar, steady state population growth rate
omega_SS     = [S,] vector, stationary steady state population distribution
surv_rate    = [S,] vector, survival rates by age
rho          = [S,] vector, mortality rates by age
g_n_vector   = [T+S,] vector, growth rate in economically active pop for each period in transition path
e            = [S,J] array, normalized effective labor units by age and ability type
------------------------------------------------------------------------
'''


def get_reduced_parameters(baseline, guid, user_modifiable, metadata):
    '''
    --------------------------------------------------------------------
    This function sets the parameters for the reduced model, which is
    simplified to run more quickly for testing.
    --------------------------------------------------------------------

    INPUTS:
    baseline        = boolean, =True if baseline tax policy, =False if reform
    guid            = string, id for reform run
    user_modifiable = boolean, =True if allow user modifiable parameters
    metadata        = boolean, =True if use metadata file for parameter
                       values (rather than what is entered in parameters below)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    read_tax_func_estimate()
    elliptical_u_est.estimation()
    read_parameter_metadata()

    OBJECTS CREATED WITHIN FUNCTION:
    See parameters defined above
    allvars = dictionary, dictionary with all parameters defined in this function

    RETURNS: allvars

    OUTPUT: None
    --------------------------------------------------------------------
    '''
    # Model Parameters
    starting_age = 40
    ending_age = 50
    S = int(ending_age-starting_age)
    J = int(1)
    T = int(2 * S)
    BW = int(10)
    lambdas = np.array([1.])
    E = int(starting_age * (S / float(ending_age - starting_age)))
    beta_annual = .96 # Carroll (JME, 2009)
    beta = beta_annual ** (float(ending_age - starting_age) / S)
    sigma = 1.5 # value from Attanasio, Banks, Meghir and Weber (JEBS, 1999)
    alpha = .35 # many use 0.33, but many find that capitals share is increasing (e.g. Elsby, Hobijn, and Sahin (BPEA, 2013))
    Z = 1.0
    delta_annual = .05 # approximately the value from Kehoe calibration exercise: http://www.econ.umn.edu/~tkehoe/classes/calibration-04.pdf
    delta = 1 - ((1 - delta_annual) ** (float(ending_age - starting_age) / S))
    ltilde = 1.0
    g_y_annual = 0.03
    g_y = (1 + g_y_annual)**(float(ending_age - starting_age) / S) - 1
    #   Ellipse parameters
    frisch = 0.4 # Frisch elasticity consistent with Altonji (JPE, 1996) and Peterman (Econ Inquiry, 2016)
    b_ellipse, upsilon = elliptical_u_est.estimation(frisch,ltilde)
    k_ellipse = 0 # this parameter is just a level shifter in utlitiy - irrelevant for analysis

    # Tax parameters:
    #   Income Tax Parameters
    ####  will call tax function estimation function here...
    ### do output such that each parameters is in a separate SxBW array

    if baseline:
        baseline_pckl = "TxFuncEst_baseline{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     baseline_pckl)
        print 'using baseline1 tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, baseline_pckl)
    else:
        policy_pckl = "TxFuncEst_policy{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     policy_pckl)
        print 'using policy1 tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, policy_pckl)


    mean_income_data = dict_params['tfunc_avginc'][0]

    etr_params = dict_params['tfunc_etr_params_S'][:S,:BW,:]
    mtrx_params = dict_params['tfunc_mtrx_params_S'][:S,:BW,:]
    mtry_params = dict_params['tfunc_mtry_params_S'][:S,:BW,:]

    # To zero out income taxes, uncomment the following 3 lines:
    # etr_params[:,:,6:] = 0.0
    # mtrx_params[:,:,6:] = 0.0
    # mtry_params[:,:,6:] = 0.0

    #   Wealth tax params
    #       These are non-calibrated values, h and m just need
    #       need to be nonzero to avoid errors. When p_wealth
    #       is zero, there is no wealth tax.
    h_wealth = 0.1
    m_wealth = 1.0
    p_wealth = 0.0
    #   Bequest and Payroll Taxes
    tau_bq = np.zeros(J)
    tau_payroll = 0.15
    retire = np.round(9.0 * S / 16.0) - 1

    # Simulation Parameters
    MINIMIZER_TOL = 1e-3
    MINIMIZER_OPTIONS = {'maxiter': 1}
    PLOT_TPI = False
    maxiter = 10
    mindist_SS = 1e-3
    mindist_TPI = 1e-3 #1e-6
    nu = .4
    flag_graphs = False
    #   Calibration parameters
    # These guesses are close to the calibrated values
    chi_b_guess = np.array([1, 100000])
    chi_n_guess = np.array([5, 6, 7, 8, 9, 10, 11, 12, 13, 14])

    # Generate Income and Demographic parameters
    omega, g_n_ss, omega_SS, surv_rate, rho, g_n_vector, imm_rates, omega_S_preTP = get_pop_objs(
        E, S, T, 1, 100, 2016, flag_graphs)
    e = np.array([[0.25, 1.25]] * 10)
    allvars = dict(locals())

    if user_modifiable:
        allvars = {k:allvars[k] for k in USER_MODIFIABLE_PARAMS}

    if metadata:
        params_meta = read_parameter_metadata()
        for k,v in allvars.iteritems():
            params_meta[k]["value"] = v
        allvars = params_meta

    return allvars


def get_full_parameters(baseline, guid, user_modifiable, metadata):
    '''
    --------------------------------------------------------------------
    This function sets the parameters for the full model.
    --------------------------------------------------------------------

    INPUTS:
    baseline        = boolean, =True if baseline tax policy, =False if reform
    guid            = string, id for reform run
    user_modifiable = boolean, =True if allow user modifiable parameters
    metadata        = boolean, =True if use metadata file for parameter
                       values (rather than what is entered in parameters below)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
    read_tax_func_estimate()
    elliptical_u_est.estimation()
    read_parameter_metadata()

    OBJECTS CREATED WITHIN FUNCTION:
    See parameters defined above
    allvars = dictionary, dictionary with all parameters defined in this function

    RETURNS: allvars

    OUTPUT: None
    --------------------------------------------------------------------
    '''
    # Model Parameters
    S = int(30) #S<30 won't meet necessary tolerances
    J = int(1)
    T = int(10 * S)
    BW = int(10)
    #lambdas = np.array([.25, .25, .2, .1, .1, .09, .01])
    #lambdas = np.array([0.5, 0.5])
    lambdas = np.array([1.,])

    starting_age = 20
    ending_age = 100
    E = int(starting_age * (S / float(ending_age - starting_age)))
    beta_annual = .96 # Carroll (JME, 2009)
    beta = beta_annual ** (float(ending_age - starting_age) / S)
    sigma = 1.5 # value from Attanasio, Banks, Meghir and Weber (JEBS, 1999)
    alpha = .35 # many use 0.33, but many find that capitals share is increasing (e.g. Elsby, Hobijn, and Sahin (BPEA, 2013))
    Z = 1.0
    delta_annual = .05 # approximately the value from Kehoe calibration exercise: http://www.econ.umn.edu/~tkehoe/classes/calibration-04.pdf
    delta = 1 - ((1 - delta_annual) ** (float(ending_age - starting_age) / S))
    ltilde = 1.0
    g_y_annual = 0.03
    g_y = (1 + g_y_annual)**(float(ending_age - starting_age) / S) - 1
    #   Ellipse parameters
    frisch = 0.4 # Frisch elasticity consistent with Altonji (JPE, 1996) and Peterman (Econ Inquiry, 2016)
    b_ellipse, upsilon = elliptical_u_est.estimation(frisch,ltilde)
    k_ellipse = 0 # this parameter is just a level shifter in utlitiy - irrelevant for analysis

    # Tax parameters:
    #   Income Tax Parameters
    #  will call tax function estimation function here...
    # do output such that each parameters is in a separate SxBW array
    # read in estimated parameters
    print 'baselines is:', baseline
    if baseline:
        baseline_pckl = "TxFuncEst_baseline{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     baseline_pckl)
        print 'using baseline tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, baseline_pckl)

    else:
        policy_pckl = "TxFuncEst_policy{}.pkl".format(guid)
        estimate_file = os.path.join(TAX_ESTIMATE_PATH,
                                     policy_pckl)
        print 'using policy tax parameters'
        dict_params = read_tax_func_estimate(estimate_file, policy_pckl)


    mean_income_data = dict_params['tfunc_avginc'][0]

    etr_params = dict_params['tfunc_etr_params_S'][:S,:BW,:]
    mtrx_params = dict_params['tfunc_mtrx_params_S'][:S,:BW,:]
    mtry_params = dict_params['tfunc_mtry_params_S'][:S,:BW,:]


    etr_params = np.zeros((S,BW,10))
    mtrx_params = np.zeros((S,BW,10))
    mtry_params = np.zeros((S,BW,10))
    etr_params[:,:,:] = dict_params['tfunc_etr_params_S'][20,0,:]
    mtrx_params[:,:,:] = dict_params['tfunc_mtrx_params_S'][20,0,:]
    mtry_params[:,:,:] = dict_params['tfunc_mtry_params_S'][20,0,:]


    mtrx_params[:,:,1] = 0.
    mtrx_params[:,:,2] = 0.
    mtrx_params[:,:,4] = 0.
    mtrx_params[:,:,8] = 0.
    mtrx_params[:,:,9] = 0.

    mtry_params[:,:,0] = 0.
    mtry_params[:,:,2] = 0.
    mtry_params[:,:,3] = 0.
    mtry_params[:,:,6] = 0.
    mtry_params[:,:,7] = 0.
    
    # # unocmmenting the block below ensures no tax rates are negative
    # etr_params[:,:,7] = 0.
    # mtrx_params[:,:,7] = 0.
    # mtry_params[:,:,7] = 0.
    # etr_params[:,:,9] = 0.
    # mtrx_params[:,:,9] = 0.
    # mtry_params[:,:,9] = 0.

    # print 'tax diffs: ', np.absolute(etr_params-etr_params.mean(axis=(0,1))).max()
    # quit()




    # set etrs and mtrs to constant rates over income/age by uncommenting following code block
    etr_params = np.zeros((S,BW,10))
    # mtrx_params = np.zeros((S,BW,10))
    # mtry_params = np.zeros((S,BW,10))
    etr_params[:,:,7] = dict_params['tfunc_avg_etr']
    # mtrx_params[:,:,7] = dict_params['tfunc_avg_mtrx']
    # mtry_params[:,:,7] = dict_params['tfunc_avg_mtry']
    # etr_params[:,:,9] = dict_params['tfunc_avg_etr']
    # mtrx_params[:,:,9] = dict_params['tfunc_avg_mtrx']
    # mtry_params[:,:,9] = dict_params['tfunc_avg_mtry']
    etr_params[:,:,5] = 1.0
    # mtrx_params[:,:,5] = 1.0
    # mtry_params[:,:,5] = 1.0

    # etr_params = np.zeros((S,BW,10))
    # mtrx_params = np.zeros((S,BW,10))
    # mtry_params = np.zeros((S,BW,10))
    # for i in xrange(S):
    #     etr_params[:,:,7] = 0.005*i
{'fraction_non_alphanumeric': 0.1962308655049581, 'fraction_numerical': 0.014537402522383749, 'mean_word_length': 1.9077612149205683, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2942446', 'n_tokens_mistral': 12404, 'n_tokens_neox': 11559, 'n_words': 3496}