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 test_runner_reform(dask_client):
p = Specifications(baseline=False, num_workers=NUM_WORKERS)
p.update_specifications(TEST_PARAM_DICT)
p.baseline_dir = BASELINE_DIR
p.output_base = REFORM_DIR
runner(p, time_path=False, client=dask_client)
