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