Ejemplo n.º 1
0
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))
Ejemplo n.º 2
0
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)