def test_SS_solver_extra(baseline, param_updates, filename, dask_client): # Test SS.SS_solver function. Provide inputs to function and # ensure that output returned matches what it has been before. p = Specifications(baseline=baseline, num_workers=NUM_WORKERS) p.update_specifications(param_updates) p.output_base = CUR_PATH b_guess = np.ones((p.S, p.J)) * 0.07 n_guess = np.ones((p.S, p.J)) * .35 * p.ltilde if p.zeta_K[-1] == 1.0: rguess = p.world_int_rate[-1] else: rguess = 0.06483431412921253 TRguess = 0.05738932081035772 factorguess = 139355.1547340256 BQguess = aggregates.get_BQ(rguess, b_guess, None, p, 'SS', False) Yguess = 0.6376591201150815 test_dict = SS.SS_solver(b_guess, n_guess, rguess, BQguess, TRguess, factorguess, Yguess, p, dask_client, False) expected_dict = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', filename)) for k, v in expected_dict.items(): print('Testing ', k) assert (np.allclose(test_dict[k], v, atol=1e-05, equal_nan=True))
def test_D_G_path(baseline_spending, Y, TR, Revenue, Gbaseline, budget_balance, expected_tuple): p = Specifications() new_param_values = { 'T': 320, 'S': 80, 'debt_ratio_ss': 1.2, 'tG1': 20, 'tG2': 256, 'alpha_T': [0.09], 'alpha_G': [0.05], 'rho_G': 0.1, 'g_y_annual': 0.03, 'baseline_spending': baseline_spending, 'budget_balance': budget_balance } p.update_specifications(new_param_values, raise_errors=False) r_gov = np.ones(p.T + p.S) * 0.03 p.g_n = np.ones(p.T + p.S) * 0.02 D0_baseline = 0.59 Gbaseline[0] = 0.05 net_revenue = Revenue pension_amount = np.zeros_like(net_revenue) UBI_outlays = np.zeros_like(net_revenue) dg_fixed_values = (Y, Revenue, pension_amount, UBI_outlays, TR, Gbaseline, D0_baseline) test_tuple = fiscal.D_G_path(r_gov, dg_fixed_values, p) for i, v in enumerate(test_tuple): assert np.allclose(v[:p.T], expected_tuple[i][:p.T])
def test_inner_loop(baseline, param_updates, filename, dask_client): # Test SS.inner_loop function. Provide inputs to function and # ensure that output returned matches what it has been before. p = Specifications(baseline=baseline, num_workers=NUM_WORKERS) p.update_specifications(param_updates) p.output_base = CUR_PATH bssmat = np.ones((p.S, p.J)) * 0.07 nssmat = np.ones((p.S, p.J)) * .4 * p.ltilde if p.zeta_K[-1] == 1.0: r = p.world_int_rate[-1] else: r = 0.05 TR = 0.12 Y = 1.3 factor = 100000 BQ = np.ones(p.J) * 0.00019646295986015257 if p.budget_balance: outer_loop_vars = (bssmat, nssmat, r, BQ, TR, factor) else: outer_loop_vars = (bssmat, nssmat, r, BQ, Y, TR, factor) test_tuple = SS.inner_loop(outer_loop_vars, p, dask_client) expected_tuple = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', filename)) for i, v in enumerate(expected_tuple): print('Max diff = ', np.absolute(test_tuple[i] - v).max()) print('Checking item = ', i) assert (np.allclose(test_tuple[i], v, atol=4e-05))
def test_sigma(sigma, dask_client): og_spec = {'frisch': 0.41, 'debt_ratio_ss': 1.0, 'sigma': sigma} p = Specifications(baseline=True, num_workers=NUM_WORKERS, baseline_dir=OUTPUT_DIR, output_base=OUTPUT_DIR) p.update_specifications(og_spec) runner(p, time_path=False, client=dask_client)
def test_update_specifications_with_dict(): spec = Specifications() new_spec_dict = { 'frisch': 0.3, } spec.update_specifications(new_spec_dict) assert spec.frisch == 0.3 assert len(spec.errors) == 0
def test_run_TPI(baseline, param_updates, filename, tmp_path, dask_client): ''' Test TPI.run_TPI function. Provide inputs to function and ensure that output returned matches what it has been before. ''' baseline_dir = os.path.join(CUR_PATH, 'baseline') if baseline: output_base = baseline_dir else: output_base = os.path.join(CUR_PATH, 'reform') p = Specifications(baseline=baseline, baseline_dir=baseline_dir, output_base=output_base, num_workers=NUM_WORKERS) test_params = TEST_PARAM_DICT.copy() test_params.update(param_updates) p.update_specifications(test_params) p.maxiter = 2 # this test runs through just two iterations # Need to run SS first to get results SS.ENFORCE_SOLUTION_CHECKS = False ss_outputs = SS.run_SS(p, client=dask_client) if p.baseline: utils.mkdirs(os.path.join(p.baseline_dir, "SS")) ss_dir = os.path.join(p.baseline_dir, "SS", "SS_vars.pkl") with open(ss_dir, "wb") as f: pickle.dump(ss_outputs, f) else: utils.mkdirs(os.path.join(p.output_base, "SS")) ss_dir = os.path.join(p.output_base, "SS", "SS_vars.pkl") with open(ss_dir, "wb") as f: pickle.dump(ss_outputs, f) TPI.ENFORCE_SOLUTION_CHECKS = False test_dict = TPI.run_TPI(p, client=dask_client) expected_dict = utils.safe_read_pickle(filename) for k, v in expected_dict.items(): print('Max diff in ', k, ' = ') try: print(np.absolute(test_dict[k][:p.T] - v[:p.T]).max()) except ValueError: print(np.absolute(test_dict[k][:p.T, :, :] - v[:p.T, :, :]).max()) for k, v in expected_dict.items(): try: assert (np.allclose(test_dict[k][:p.T], v[:p.T], rtol=1e-04, atol=1e-04)) except ValueError: assert (np.allclose(test_dict[k][:p.T, :, :], v[:p.T, :, :], rtol=1e-04, atol=1e-04))
def test_implement_reform(): specs = Specifications() new_specs = {'tG1': 30, 'T': 80, 'frisch': 0.3, 'tax_func_type': 'DEP'} specs.update_specifications(new_specs) assert specs.frisch == 0.3 assert specs.tG1 == 30 assert specs.T == 80 assert specs.tax_func_type == 'DEP' assert len(specs.errors) == 0
def test_update_specification_with_json(): spec = Specifications() new_spec_json = """ { "frisch": 0.3 } """ spec.update_specifications(new_spec_json) assert spec.frisch == 0.3 assert len(spec.errors) == 0
def test_implement_bad_reform2(): specs = Specifications() # tG1 has an upper bound at T / 2 new_specs = {'T': 80, 'tax_func_type': 'not_a_functional_form'} specs.update_specifications(new_specs, raise_errors=False) assert len(specs.errors) > 0 assert specs.errors['tax_func_type'][0] == ( 'tax_func_type "not_a_functional_form" must be in list of ' + 'choices DEP, DEP_totalinc, GS, linear.')
def test_implement_bad_reform1(): specs = Specifications() # tG1 has an upper bound at T / 2 new_specs = { 'tG1': 50, 'T': 80, } specs.update_specifications(new_specs, raise_errors=False) assert len(specs.errors) == 0
def test_run_small(time_path, dask_client): from ogcore.execute import runner # Monkey patch enforcement flag since small data won't pass checks SS.ENFORCE_SOLUTION_CHECKS = False TPI.ENFORCE_SOLUTION_CHECKS = False SS.MINIMIZER_TOL = 1e-6 TPI.MINIMIZER_TOL = 1e-6 p = Specifications(baseline=True, num_workers=NUM_WORKERS, baseline_dir=OUTPUT_DIR, output_base=OUTPUT_DIR) p.update_specifications(TEST_PARAM_DICT) runner(p, time_path=time_path, client=dask_client)
def test_get_D_ss(budget_balance, expected_tuple): ''' Test of the fiscla.get_D_ss() function. ''' r_gov = 0.03 Y = 1.176255339 p = Specifications() p.debt_ratio_ss = 1.2 p.budget_balance = budget_balance p.g_n_ss = 0.02 test_tuple = fiscal.get_D_ss(r_gov, Y, p) for i, v in enumerate(test_tuple): assert np.allclose(v, expected_tuple[i])
def test_get_biz_tax(): # Test function for business tax receipts p = Specifications() new_param_values = { 'cit_rate': [0.20], 'delta_tau_annual': [0.06] } p.update_specifications(new_param_values) p.T = 3 w = np.array([1.2, 1.1, 1.2]) Y = np.array([3.0, 7.0, 3.0]) L = np.array([2.0, 3.0, 2.0]) K = np.array([5.0, 6.0, 5.0]) biz_tax = tax.get_biz_tax(w, Y, L, K, p, 'TPI') assert np.allclose(biz_tax, np.array([0.0102, 0.11356, 0.0102]))
def test_get_G_ss(budget_balance, expected_G): ''' Test of the fiscla.get_G_ss() function. ''' Y = 2.2 net_revenue = 2.3 pension_amount = 0.0 TR = 1.6 UBI = 0.0 new_borrowing = 0.072076633 debt_service = 0.042345192 p = Specifications() p.budget_balance = budget_balance test_G = fiscal.get_G_ss(Y, net_revenue, pension_amount, TR, UBI, new_borrowing, debt_service, p) assert np.allclose(test_G, expected_G)
def test_get_y(): ''' Test of household.get_y() function. ''' r_hh = np.array([0.05, 0.04, 0.09]) w = np.array([1.2, 0.8, 2.5]) b_s = np.array([0.5, 0.99, 9]) n = np.array([0.8, 3.2, 0.2]) expected_y = np.array([0.9754, 3.8796, 0.91]) p = Specifications() # p.update_specifications({'S': 4, 'J': 1}) p.S = 3 p.e = np.array([0.99, 1.5, 0.2]) test_y = household.get_y(r_hh, w, b_s, n, p) assert np.allclose(test_y, expected_y)
def validate_inputs(meta_param_dict, adjustment, errors_warnings): # ogusa doesn't look at meta_param_dict for validating inputs. params = Specifications() params.adjust(adjustment["OG-USA Parameters"], raise_errors=False) errors_warnings["OG-USA Parameters"]["errors"].update(params.errors) # Validate TC parameter inputs pol_params = {} # drop checkbox parameters. for param, data in list(adjustment["Tax-Calculator Parameters"].items()): if not param.endswith("checkbox"): pol_params[param] = data iit_params = TCParams() iit_params.adjust(pol_params, raise_errors=False) errors_warnings["Tax-Calculator Parameters"]["errors"].update( iit_params.errors ) return {"errors_warnings": errors_warnings}
def test_resource_constraint(): """ Test resource constraint equation. """ p = Specifications() p.delta = 0.05 Y = np.array([48, 55, 2, 99, 8]) C = np.array([33, 44, 0.4, 55, 6]) G = np.array([4, 5, 0.01, 22, 0]) I = np.array([20, 5, 0.6, 10, 1]) K_f = np.array([0, 0, 0.2, 3, 0.05]) new_borrowing_f = np.array([0, 0.1, 0.3, 4, 0.5]) debt_service_f = np.array([0.1, 0.1, 0.3, 2, 0.02]) r = np.array([0.03, 0.04, 0.03, 0.06, 0.01]) expected = np.array([-9.1, 1, 0.974, 13.67, 1.477]) test_RC = aggr.resource_constraint(Y, C, G, I, K_f, new_borrowing_f, debt_service_f, r, p) assert (np.allclose(test_RC, expected))
def test_get_initial_SS_values(baseline, param_updates, filename, dask_client): p = Specifications(baseline=baseline, num_workers=NUM_WORKERS) p.update_specifications(param_updates) p.baseline_dir = os.path.join(CUR_PATH, 'test_io_data', 'OUTPUT') p.output_base = os.path.join(CUR_PATH, 'test_io_data', 'OUTPUT') test_tuple = TPI.get_initial_SS_values(p) (test_initial_values, test_ss_vars, test_theta, test_baseline_values) = test_tuple expected_tuple = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', filename)) (exp_initial_values, exp_ss_vars, exp_theta, exp_baseline_values) = expected_tuple (B0, b_sinit, b_splus1init, factor, initial_b, initial_n) = exp_initial_values B0 = aggr.get_B(exp_ss_vars['bssmat_splus1'], p, 'SS', True) initial_b = (exp_ss_vars['bssmat_splus1'] * (exp_ss_vars['Bss'] / B0)) B0 = aggr.get_B(initial_b, p, 'SS', True) b_sinit = np.array( list(np.zeros(p.J).reshape(1, p.J)) + list(initial_b[:-1])) b_splus1init = initial_b exp_initial_values = (B0, b_sinit, b_splus1init, factor, initial_b, initial_n) for i, v in enumerate(exp_initial_values): assert (np.allclose(test_initial_values[i], v, equal_nan=True)) if p.baseline_spending: for i, v in enumerate(exp_baseline_values): assert (np.allclose(test_baseline_values[i], v, equal_nan=True)) assert (np.allclose(test_theta, exp_theta)) for k, v in exp_ss_vars.items(): assert (np.allclose(test_ss_vars[k], v, equal_nan=True))
def run_micro_macro(og_spec, guid, client): guid = '' start_time = time.time() REFORM_DIR = os.path.join(CUR_PATH, "OUTPUT_REFORM_" + guid) BASELINE_DIR = os.path.join(CUR_PATH, "OUTPUT_BASELINE" + guid) with open("log_{}.log".format(guid), 'w') as f: f.write("guid: {}\n".format(guid)) f.write("og_spec: {}\n".format(og_spec)) ''' ------------------------------------------------------------------------ Run baseline ------------------------------------------------------------------------ ''' p = Specifications( baseline=True, num_workers=NUM_WORKERS, baseline_dir=BASELINE_DIR, output_base=BASELINE_DIR) p.update_specifications(og_spec) runner(p, time_path=True, client=client) ''' ------------------------------------------------------------------------ Run reform ------------------------------------------------------------------------ ''' p = Specifications( baseline=False, num_workers=NUM_WORKERS, baseline_dir=BASELINE_DIR, output_base=REFORM_DIR) p.update_specifications(og_spec) runner(p, time_path=True, client=client) time.sleep(0.5) base_tpi = safe_read_pickle( os.path.join(BASELINE_DIR, 'TPI', 'TPI_vars.pkl')) base_params = safe_read_pickle( os.path.join(BASELINE_DIR, 'model_params.pkl')) reform_tpi = safe_read_pickle( os.path.join(REFORM_DIR, 'TPI', 'TPI_vars.pkl')) reform_params = safe_read_pickle( os.path.join(REFORM_DIR, 'model_params.pkl')) ans = ot.macro_table( base_tpi, base_params, reform_tpi=reform_tpi, reform_params=reform_params, var_list=['Y', 'C', 'K', 'L', 'r', 'w'], output_type='pct_diff', num_years=10, start_year=base_params.start_year) print("total time was ", (time.time() - start_time)) return ans
def test_get_TR(baseline, budget_balance, baseline_spending, method, expected_TR): ''' Test of the fiscal.get_TR() function. ''' Y = 3.2 TR = 1.5 G = 0.0 agg_pension_outlays = 0.0 UBI_outlays = 0.0 total_tax_revenue = 1.9 p = Specifications(baseline=baseline) p.budget_balance = budget_balance p.baseline_spending = baseline_spending if method == 'TPI': Y = np.ones(p.T * p.S) * Y TR = np.ones(p.T * p.S) * TR total_tax_revenue = np.ones(p.T * p.S) * total_tax_revenue test_TR = fiscal.get_TR(Y, TR, G, total_tax_revenue, agg_pension_outlays, UBI_outlays, p, method) assert np.allclose(test_TR, expected_TR)
def test_twist_doughnut(file_inputs, file_outputs): ''' Test TPI.twist_doughnut function. Provide inputs to function and ensure that output returned matches what it has been before. ''' input_tuple = utils.safe_read_pickle(file_inputs) (guesses, r, w, bq, tr, theta, factor, ubi, j, s, t, tau_c, etr_params, mtrx_params, mtry_params, initial_b) = input_tuple p = Specifications() input_tuple = (guesses, r, w, bq, tr, theta, factor, ubi, j, s, t, tau_c, etr_params, mtrx_params, mtry_params, initial_b, p) test_list = TPI.twist_doughnut(*input_tuple) expected_list = utils.safe_read_pickle(file_outputs) assert (np.allclose(np.array(test_list), np.array(expected_list)))
def get_inputs(meta_param_dict): meta_params = MetaParams() meta_params.adjust(meta_param_dict) # Set default OG-USA parameters ogusa_params = Specifications() ogusa_params.update_specifications( json.load( open( os.path.join( "..", "..", "ogusa", "ogusa_default_parameters.json" ) ) ) ) ogusa_params.start_year = meta_params.year filtered_ogusa_params = OrderedDict() filter_list = [ "chi_n_80", "chi_b", "eta", "zeta", "constant_demographics", "ltilde", "use_zeta", "constant_rates", "zero_taxes", "analytical_mtrs", "age_specific", "gamma", "epsilon", "start_year", ] for k, v in ogusa_params.dump().items(): if ( (k not in filter_list) and (v.get("section_1", False) != "Model Solution Parameters") and (v.get("section_2", False) != "Model Dimensions") ): filtered_ogusa_params[k] = v print("filtered ogusa = ", k) # Set default TC params iit_params = TCParams() iit_params.set_state(year=meta_params.year.tolist()) filtered_iit_params = OrderedDict() for k, v in iit_params.dump().items(): if k == "schema" or v.get("section_1", False): filtered_iit_params[k] = v default_params = { "OG-USA Parameters": filtered_ogusa_params, "Tax-Calculator Parameters": filtered_iit_params, } return { "meta_parameters": meta_params.dump(), "model_parameters": default_params, }
def test_inner_loop(): # Test TPI.inner_loop function. Provide inputs to function and # ensure that output returned matches what it has been before. input_tuple = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', 'tpi_inner_loop_inputs.pkl')) guesses, outer_loop_vars, initial_values, ubi, j, ind = input_tuple p = Specifications() test_tuple = TPI.inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p) expected_tuple = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', 'tpi_inner_loop_outputs.pkl')) for i, v in enumerate(expected_tuple): assert (np.allclose(test_tuple[i], v))
def test_firstdoughnutring(): # Test TPI.firstdoughnutring function. Provide inputs to function and # ensure that output returned matches what it has been before. input_tuple = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', 'firstdoughnutring_inputs.pkl')) guesses, r, w, bq, tr, theta, factor, ubi, j, initial_b = input_tuple p = Specifications() test_list = TPI.firstdoughnutring(guesses, r, w, bq, tr, theta, factor, ubi, j, initial_b, p) expected_list = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', 'firstdoughnutring_outputs.pkl')) assert (np.allclose(np.array(test_list), np.array(expected_list)))
def test_constant_demographics_TPI(dask_client): ''' This tests solves the model under the assumption of constant demographics, a balanced budget, and tax functions that do not vary over time. In this case, given how initial guesses for the time path are made, the time path should be solved for on the first iteration and the values all along the time path should equal their steady-state values. ''' # Create output directory structure spec = Specifications(output_base=CUR_PATH, baseline_dir=OUTPUT_DIR, baseline=True, num_workers=NUM_WORKERS) og_spec = { 'constant_demographics': True, 'budget_balance': True, 'zero_taxes': True, 'maxiter': 2, 'r_gov_shift': 0.0, 'zeta_D': [0.0, 0.0], 'zeta_K': [0.0, 0.0], 'debt_ratio_ss': 1.0, 'initial_foreign_debt_ratio': 0.0, 'start_year': 2019, 'cit_rate': [0.0], 'PIA_rate_bkt_1': 0.0, 'PIA_rate_bkt_2': 0.0, 'PIA_rate_bkt_3': 0.0, 'eta': (spec.omega_SS.reshape(spec.S, 1) * spec.lambdas.reshape(1, spec.J)) } spec.update_specifications(og_spec) spec.etr_params = np.zeros( (spec.T + spec.S, spec.S, spec.etr_params.shape[2])) spec.mtrx_params = np.zeros( (spec.T + spec.S, spec.S, spec.mtrx_params.shape[2])) spec.mtry_params = np.zeros( (spec.T + spec.S, spec.S, spec.mtry_params.shape[2])) # Run SS ss_outputs = SS.run_SS(spec, client=dask_client) # save SS results utils.mkdirs(os.path.join(OUTPUT_DIR, "SS")) ss_dir = os.path.join(OUTPUT_DIR, "SS", "SS_vars.pkl") with open(ss_dir, "wb") as f: pickle.dump(ss_outputs, f) # Run TPI tpi_output = TPI.run_TPI(spec, client=dask_client) assert (np.allclose(tpi_output['bmat_splus1'][:spec.T, :, :], ss_outputs['bssmat_splus1']))
def run_model(meta_param_dict, adjustment): """ Initializes classes from OG-USA that compute the model under different policies. Then calls function get output objects. """ print("Meta_param_dict = ", meta_param_dict) print("adjustment dict = ", adjustment) meta_params = MetaParams() meta_params.adjust(meta_param_dict) if meta_params.data_source == "PUF": data = retrieve_puf(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) # set name of cached baseline file in case use below cached_pickle = "TxFuncEst_baseline_PUF.pkl" else: data = "cps" # set name of cached baseline file in case use below cached_pickle = "TxFuncEst_baseline_CPS.pkl" # Get TC params adjustments iit_mods = convert_policy_adjustment( adjustment["Tax-Calculator Parameters"] ) # Create output directory structure base_dir = os.path.join(CUR_DIR, BASELINE_DIR) reform_dir = os.path.join(CUR_DIR, REFORM_DIR) dirs = [base_dir, reform_dir] for _dir in dirs: utils.mkdirs(_dir) # Dask parmeters client = Client() num_workers = 5 # TODO: Swap to these parameters when able to specify tax function # and model workers separately # num_workers_txf = 5 # num_workers_mod = 6 # whether to estimate tax functions from microdata run_micro = True time_path = meta_param_dict["time_path"][0]["value"] # filter out OG-USA params that will not change between baseline and # reform runs (these are the non-policy parameters) filtered_ogusa_params = {} constant_param_set = { "frisch", "beta_annual", "sigma", "g_y_annual", "gamma", "epsilon", "Z", "delta_annual", "small_open", "world_int_rate", "initial_debt_ratio", "initial_foreign_debt_ratio", "zeta_D", "zeta_K", "tG1", "tG2", "rho_G", "debt_ratio_ss", "budget_balance", } filtered_ogusa_params = OrderedDict() for k, v in adjustment["OG-USA Parameters"].items(): if k in constant_param_set: filtered_ogusa_params[k] = v # Solve baseline model start_year = meta_param_dict["year"][0]["value"] if start_year == 2020: OGPATH = inspect.getfile(SS) OGDIR = os.path.dirname(OGPATH) tax_func_path = None # os.path.join(OGDIR, 'data', 'tax_functions', # cached_pickle) run_micro_baseline = True else: tax_func_path = None run_micro_baseline = True base_spec = { **{ "start_year": start_year, "tax_func_type": "DEP", "age_specific": False, }, **filtered_ogusa_params, } base_params = Specifications( output_base=base_dir, baseline_dir=base_dir, baseline=True, num_workers=num_workers, ) base_params.update_specifications( json.load( open( os.path.join( "..", "..", "ogusa", "ogusa_default_parameters.json" ) ) ) ) base_params.update_specifications(base_spec) BW = TC_LAST_YEAR - start_year + 1 base_params.BW = BW # Will need to figure out how to handle default tax functions here # For now, estimating tax functions even for baseline c_base = Calibration( base_params, iit_reform={}, estimate_tax_functions=True, data=data, client=client, ) # update tax function parameters in Specifications Object d_base = c_base.get_dict() # additional parameters to change updated_txfunc_params = { "etr_params": d_base["etr_params"], "mtrx_params": d_base["mtrx_params"], "mtry_params": d_base["mtry_params"], "mean_income_data": d_base["mean_income_data"], "frac_tax_payroll": d_base["frac_tax_payroll"], } base_params.update_specifications(updated_txfunc_params) base_ss = SS.run_SS(base_params, client=client) utils.mkdirs(os.path.join(base_dir, "SS")) base_ss_dir = os.path.join(base_dir, "SS", "SS_vars.pkl") with open(base_ss_dir, "wb") as f: pickle.dump(base_ss, f) if time_path: base_tpi = TPI.run_TPI(base_params, client=client) tpi_dir = os.path.join(base_dir, "TPI", "TPI_vars.pkl") with open(tpi_dir, "wb") as f: pickle.dump(base_tpi, f) else: base_tpi = None # Solve reform model reform_spec = base_spec reform_spec.update(adjustment["OG-USA Parameters"]) reform_params = Specifications( output_base=reform_dir, baseline_dir=base_dir, baseline=False, num_workers=num_workers, ) reform_params.update_specifications( json.load( open( os.path.join( "..", "..", "ogusa", "ogusa_default_parameters.json" ) ) ) ) reform_params.update_specifications(reform_spec) reform_params.BW = BW c_reform = Calibration( reform_params, iit_reform=iit_mods, estimate_tax_functions=True, data=data, client=client, ) # update tax function parameters in Specifications Object d_reform = c_reform.get_dict() # additional parameters to change updated_txfunc_params = { "etr_params": d_reform["etr_params"], "mtrx_params": d_reform["mtrx_params"], "mtry_params": d_reform["mtry_params"], "mean_income_data": d_reform["mean_income_data"], "frac_tax_payroll": d_reform["frac_tax_payroll"], } reform_params.update_specifications(updated_txfunc_params) reform_ss = SS.run_SS(reform_params, client=client) utils.mkdirs(os.path.join(reform_dir, "SS")) reform_ss_dir = os.path.join(reform_dir, "SS", "SS_vars.pkl") with open(reform_ss_dir, "wb") as f: pickle.dump(reform_ss, f) if time_path: reform_tpi = TPI.run_TPI(reform_params, client=client) else: reform_tpi = None comp_dict = comp_output( base_params, base_ss, reform_params, reform_ss, time_path, base_tpi, reform_tpi, ) # Shut down client and make sure all of its references are # cleaned up. client.close() del client return comp_dict
''' ineq = Inequality(dist, pop_weights, ability_weights, S, J) top_share = ineq.top_share(0.05) assert np.allclose(top_share, 0.285714286) def test_to_timepath_shape(): ''' Test of function that converts vector to time path conformable array ''' in_array = np.ones(40) test_array = utils.to_timepath_shape(in_array) assert test_array.shape == (40, 1, 1) p = Specifications() p.T = 40 p.S = 3 p.J = 1 x1 = np.ones((p.S, p.J)) * 0.4 xT = np.ones((p.S, p.J)) * 5.0 expected1 = np.tile( np.array([ 0.4, 0.51794872, 0.63589744, 0.75384615, 0.87179487, 0.98974359, 1.10769231, 1.22564103, 1.34358974, 1.46153846, 1.57948718, 1.6974359, 1.81538462, 1.93333333, 2.05128205, 2.16923077, 2.28717949, 2.40512821, 2.52307692, 2.64102564, 2.75897436, 2.87692308, 2.99487179, 3.11282051, 3.23076923, 3.34871795, 3.46666667, 3.58461538, 3.7025641, 3.82051282, 3.93846154, 4.05641026, 4.17435897, 4.29230769, 4.41025641, 4.52820513, 4.64615385, 4.76410256, 4.88205128, 5., 5.0, 5.0, 5.0 ]).reshape(p.T + p.S, 1, 1), (1, p.S, p.J))
def test_run_SS(baseline, param_updates, filename, dask_client): # Test SS.run_SS function. Provide inputs to function and # ensure that output returned matches what it has been before. SS.ENFORCE_SOLUTION_CHECKS = True #False # if running reform, then need to solve baseline first to get values if baseline is False: p_base = Specifications(output_base=constants.BASELINE_DIR, baseline_dir=constants.BASELINE_DIR, baseline=True, num_workers=NUM_WORKERS) p_base.update_specifications(param_updates) if p_base.use_zeta: p_base.update_specifications({ 'initial_guess_r_SS': 0.10, 'initial_guess_TR_SS': 0.02 }) p_base.baseline_spending = False base_ss_outputs = SS.run_SS(p_base, client=dask_client) utils.mkdirs(os.path.join(constants.BASELINE_DIR, "SS")) ss_dir = os.path.join(constants.BASELINE_DIR, "SS", "SS_vars.pkl") with open(ss_dir, "wb") as f: pickle.dump(base_ss_outputs, f) # now run specification for test p = Specifications(baseline=baseline, num_workers=NUM_WORKERS) p.update_specifications(param_updates) test_dict = SS.run_SS(p, client=dask_client) expected_dict = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', filename)) for k, v in expected_dict.items(): print('Checking item = ', k) assert (np.allclose(test_dict[k], v, atol=1e-06))
def main(): # Define parameters to use for multiprocessing client = Client() num_workers = min(multiprocessing.cpu_count(), 7) print("Number of workers = ", num_workers) # Directories to save data CUR_DIR = os.path.dirname(os.path.realpath(__file__)) base_dir = os.path.join(CUR_DIR, "OG-USA-Example", "OUTPUT_BASELINE") reform_dir = os.path.join(CUR_DIR, "OG-USA-Example", "OUTPUT_REFORM") """ ------------------------------------------------------------------------ Run baseline policy ------------------------------------------------------------------------ """ # Set up baseline parameterization p = Specifications( baseline=True, num_workers=num_workers, baseline_dir=base_dir, output_base=base_dir, ) # Update parameters for baseline from default json file p.update_specifications( json.load( open( os.path.join(CUR_DIR, "..", "ogusa", "ogusa_default_parameters.json")))) # Run model start_time = time.time() runner(p, time_path=True, client=client) print("run time = ", time.time() - start_time) """ ------------------------------------------------------------------------ Run reform policy ------------------------------------------------------------------------ """ # Grab a reform JSON file already in Tax-Calculator # In this example the 'reform' is a change to 2017 law (the # baseline policy is tax law in 2018) reform_url = ("github://*****:*****@main/psl_examples/" + "taxcalc/2017_law.json") ref = Calculator.read_json_param_objects(reform_url, None) iit_reform = ref["policy"] # create new Specifications object for reform simulation p2 = Specifications( baseline=False, num_workers=num_workers, baseline_dir=base_dir, output_base=reform_dir, ) # Update parameters for baseline from default json file p2.update_specifications( json.load( open( os.path.join(CUR_DIR, "..", "ogusa", "ogusa_default_parameters.json")))) # Use calibration class to estimate reform tax functions from # Tax-Calculator, specifing reform for Tax-Calculator in iit_reform c2 = Calibration(p2, iit_reform=iit_reform, estimate_tax_functions=True, client=client) # update tax function parameters in Specifications Object d = c2.get_dict() # additional parameters to change updated_params = { "cit_rate": [0.35], "etr_params": d["etr_params"], "mtrx_params": d["mtrx_params"], "mtry_params": d["mtry_params"], "mean_income_data": d["mean_income_data"], "frac_tax_payroll": d["frac_tax_payroll"], } p2.update_specifications(updated_params) # Run model start_time = time.time() runner(p2, time_path=True, client=client) print("run time = ", time.time() - start_time) client.close() """ ------------------------------------------------------------------------ Save some results of simulations ------------------------------------------------------------------------ """ base_tpi = safe_read_pickle(os.path.join(base_dir, "TPI", "TPI_vars.pkl")) base_params = safe_read_pickle(os.path.join(base_dir, "model_params.pkl")) reform_tpi = safe_read_pickle( os.path.join(reform_dir, "TPI", "TPI_vars.pkl")) reform_params = safe_read_pickle( os.path.join(reform_dir, "model_params.pkl")) ans = ot.macro_table( base_tpi, base_params, reform_tpi=reform_tpi, reform_params=reform_params, var_list=["Y", "C", "K", "L", "r", "w"], output_type="pct_diff", num_years=10, start_year=base_params.start_year, ) # create plots of output op.plot_all(base_dir, reform_dir, os.path.join(CUR_DIR, "OG-USA_example_plots")) print("Percentage changes in aggregates:", ans) # save percentage change output to csv file ans.to_csv("ogusa_example_output.csv")
@pytest.fixture(scope="module") def dask_client(): cluster = LocalCluster(n_workers=NUM_WORKERS, threads_per_worker=2) client = Client(cluster) yield client # teardown client.close() cluster.close() input_tuple = utils.safe_read_pickle( os.path.join(CUR_PATH, 'test_io_data', 'SS_fsolve_inputs.pkl')) (bssmat, nssmat, TR_ss, factor_ss) = input_tuple # Parameterize the baseline, closed econ case p1 = Specifications(baseline=True) p1.update_specifications({'zeta_D': [0.0], 'zeta_K': [0.0]}) guesses1 = np.array( [0.06, 0.016, 0.02, 0.02, 0.01, 0.01, 0.02, 0.003, -0.07, 0.051]) args1 = (bssmat, nssmat, None, None, p1, None) expected1 = np.array([ -0.026632037158481975, -0.0022739752626707334, -0.01871875707724979, -0.01791935965422934, 0.005996289165268601, 0.00964100151012603, -0.01953460990186908, -0.0029633389016814967, 0.1306862551496613, 0.11574464544202477 ]) # Parameterize the reform, closed econ case p2 = Specifications(baseline=False) p2.update_specifications({'zeta_D': [0.0], 'zeta_K': [0.0]}) guesses2 = np.array([0.06, 0.016, 0.02, 0.02, 0.01, 0.01, 0.02, 0.003, -0.07]) args2 = (bssmat, nssmat, None, 0.51, p2, None)