Esempio n. 1
0
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))
Esempio n. 2
0
def inner_loop(outer_loop_vars, p, client):
    '''
    This function solves for the inner loop of the SS.  That is, given
    the guesses of the outer loop variables (r, w, TR, factor) this
    function solves the households' problems in the SS.

    Args:
        outer_loop_vars (tuple): tuple of outer loop variables,
            (bssmat, nssmat, r, BQ, TR, factor) or
            (bssmat, nssmat, r, BQ, Y, TR, factor)
        bssmat (Numpy array): initial guess at savings, size = SxJ
        nssmat (Numpy array): initial guess at labor supply, size = SxJ
        BQ (array_like): aggregate bequest amount(s)
        Y (scalar): real GDP
        TR (scalar): lump sum transfer amount
        factor (scalar): scaling factor converting model units to dollars
        w (scalar): real wage rate
        p (OG-Core Specifications object): model parameters
        client (Dask client object): client

    Returns:
        (tuple): results from household solution:

            * euler_errors (Numpy array): errors terms from FOCs,
                size = 2SxJ
            * bssmat (Numpy array): savings, size = SxJ
            * nssmat (Numpy array): labor supply, size = SxJ
            * new_r (scalar): real interest rate on firm capital
            * new_r_gov (scalar): real interest rate on government debt
            * new_r_p (scalar): real interest rate on household
                portfolio
            * new_w (scalar): real wage rate
            * new_TR (scalar): lump sum transfer amount
            * new_Y (scalar): real GDP
            * new_factor (scalar): scaling factor converting model
                units to dollars
            * new_BQ (array_like): aggregate bequest amount(s)
            * average_income_model (scalar): average income in model
                units

    '''
    # unpack variables to pass to function
    if p.budget_balance:
        bssmat, nssmat, r, BQ, TR, factor = outer_loop_vars
        r_p = r
        Y = 1.0  # placeholder
        K = 1.0  # placeholder
    else:
        bssmat, nssmat, r, BQ, Y, TR, factor = outer_loop_vars
        K = firm.get_K_from_Y(Y, r, p, 'SS')
    # initialize array for euler errors
    euler_errors = np.zeros((2 * p.S, p.J))

    w = firm.get_w_from_r(r, p, 'SS')
    r_gov = fiscal.get_r_gov(r, p)
    D, D_d, D_f, new_borrowing, debt_service, new_borrowing_f =\
        fiscal.get_D_ss(r_gov, Y, p)
    r_p = aggr.get_r_p(r, r_gov, K, D)
    bq = household.get_bq(BQ, None, p, 'SS')
    tr = household.get_tr(TR, None, p, 'SS')
    ubi = p.ubi_nom_array[-1, :, :] / factor

    lazy_values = []
    for j in range(p.J):
        guesses = np.append(bssmat[:, j], nssmat[:, j])
        euler_params = (r_p, w, bq[:, j], tr[:, j], ubi[:, j], factor, j, p)
        lazy_values.append(
            delayed(opt.fsolve)(euler_equation_solver,
                                guesses * .9,
                                args=euler_params,
                                xtol=MINIMIZER_TOL,
                                full_output=True))
    if client:
        futures = client.compute(lazy_values, num_workers=p.num_workers)
        results = client.gather(futures)
    else:
        results = results = compute(*lazy_values,
                                    scheduler=dask.multiprocessing.get,
                                    num_workers=p.num_workers)

    # for j, result in results.items():
    for j, result in enumerate(results):
        [solutions, infodict, ier, message] = result
        euler_errors[:, j] = infodict['fvec']
        bssmat[:, j] = solutions[:p.S]
        nssmat[:, j] = solutions[p.S:]

    L = aggr.get_L(nssmat, p, 'SS')
    B = aggr.get_B(bssmat, p, 'SS', False)
    K_demand_open = firm.get_K(L, p.world_int_rate[-1], p, 'SS')
    K, K_d, K_f = aggr.get_K_splits(B, K_demand_open, D_d, p.zeta_K[-1])
    Y = firm.get_Y(K, L, p, 'SS')
    if p.zeta_K[-1] == 1.0:
        new_r = p.world_int_rate[-1]
    else:
        new_r = firm.get_r(Y, K, p, 'SS')
    new_w = firm.get_w_from_r(new_r, p, 'SS')

    b_s = np.array(list(np.zeros(p.J).reshape(1, p.J)) + list(bssmat[:-1, :]))
    new_r_gov = fiscal.get_r_gov(new_r, p)
    new_r_p = aggr.get_r_p(new_r, new_r_gov, K, D)
    average_income_model = ((new_r_p * b_s + new_w * p.e * nssmat) *
                            p.omega_SS.reshape(p.S, 1) *
                            p.lambdas.reshape(1, p.J)).sum()
    if p.baseline:
        new_factor = p.mean_income_data / average_income_model
    else:
        new_factor = factor
    new_BQ = aggr.get_BQ(new_r_p, bssmat, None, p, 'SS', False)
    new_bq = household.get_bq(new_BQ, None, p, 'SS')
    tr = household.get_tr(TR, None, p, 'SS')
    theta = tax.replacement_rate_vals(nssmat, new_w, new_factor, None, p)
    etr_params_3D = np.tile(
        np.reshape(p.etr_params[-1, :, :], (p.S, 1, p.etr_params.shape[2])),
        (1, p.J, 1))
    taxss = tax.net_taxes(new_r_p, new_w, b_s, nssmat, new_bq, factor, tr, ubi,
                          theta, None, None, False, 'SS', p.e, etr_params_3D,
                          p)
    cssmat = household.get_cons(new_r_p, new_w, b_s, bssmat, nssmat, new_bq,
                                taxss, p.e, p.tau_c[-1, :, :], p)
    total_tax_revenue, _, agg_pension_outlays, UBI_outlays, _, _, _, _, _, _ =\
        aggr.revenue(new_r_p, new_w, b_s, nssmat, new_bq, cssmat, Y, L,
                     K, factor, ubi, theta, etr_params_3D, p, 'SS')
    G = fiscal.get_G_ss(Y, total_tax_revenue, agg_pension_outlays, TR,
                        UBI_outlays, new_borrowing, debt_service, p)
    new_TR = fiscal.get_TR(Y, TR, G, total_tax_revenue, agg_pension_outlays,
                           UBI_outlays, p, 'SS')

    return euler_errors, bssmat, nssmat, new_r, new_r_gov, new_r_p, \
        new_w, new_TR, Y, new_factor, new_BQ, average_income_model
Esempio n. 3
0
def test_get_BQ(r, b_splus1, j, p, method, PreTP, expected):
    """
    Test of aggregate bequest function.
    """
    BQ = aggr.get_BQ(r, b_splus1, j, p, method, PreTP)
    assert np.allclose(BQ, expected)
Esempio n. 4
0
def run_SS(p, client=None):
    '''
    Solve for steady-state equilibrium of OG-Core.

    Args:
        p (OG-Core Specifications object): model parameters
        client (Dask client object): client

    Returns:
        output (dictionary): dictionary with steady-state solution
            results

    '''
    # For initial guesses of w, r, TR, and factor, we use values that
    # are close to some steady state values.
    if p.baseline:
        if p.zeta_K[-1] == 1.0:
            rguess = p.world_int_rate[-1]
        else:
            rguess = p.initial_guess_r_SS
        if p.use_zeta:
            b_guess = np.ones((p.S, p.J)) * 0.0055
            n_guess = np.ones((p.S, p.J)) * .4 * p.ltilde
        else:
            b_guess = np.ones((p.S, p.J)) * 0.07
            n_guess = np.ones((p.S, p.J)) * .35 * p.ltilde
        TRguess = p.initial_guess_TR_SS
        factorguess = p.initial_guess_factor_SS
        BQguess = aggr.get_BQ(rguess, b_guess, None, p, 'SS', False)
        ss_params_baseline = (b_guess, n_guess, None, None, p, client)
        if p.use_zeta:
            BQguess = 0.12231465279007188
            guesses = [rguess] + list([BQguess]) + [TRguess, factorguess]
        else:
            guesses = [rguess] + list(BQguess) + [TRguess, factorguess]
        [solutions_fsolve, infodict, ier, message] =\
            opt.fsolve(SS_fsolve, guesses, args=ss_params_baseline,
                       xtol=p.mindist_SS, full_output=True)
        if ENFORCE_SOLUTION_CHECKS and not ier == 1:
            raise RuntimeError('Steady state equilibrium not found')
        rss = solutions_fsolve[0]
        BQss = solutions_fsolve[1:-2]
        TR_ss = solutions_fsolve[-2]
        factor_ss = solutions_fsolve[-1]
        Yss = TR_ss / p.alpha_T[-1]  # may not be right - if budget_balance
        # = True, but that's ok - will be fixed in SS_solver
        fsolve_flag = True
        # Return SS values of variables
        output = SS_solver(b_guess, n_guess, rss, BQss, TR_ss, factor_ss, Yss,
                           p, client, fsolve_flag)
    else:
        # Use the baseline solution to get starting values for the reform
        baseline_ss_dir = os.path.join(p.baseline_dir, 'SS', 'SS_vars.pkl')
        ss_solutions = utils.safe_read_pickle(baseline_ss_dir)
        # use baseline solution as starting values if dimensions match
        if ss_solutions['bssmat_splus1'].shape == (p.S, p.J):
            (b_guess, n_guess, rguess, BQguess, TRguess, Yguess,
             factor) =\
                (ss_solutions['bssmat_splus1'], ss_solutions['nssmat'],
                 ss_solutions['rss'], ss_solutions['BQss'],
                 ss_solutions['TR_ss'], ss_solutions['Yss'],
                 ss_solutions['factor_ss'])
        else:
            if p.use_zeta:
                b_guess = np.ones((p.S, p.J)) * 0.0055
                n_guess = np.ones((p.S, p.J)) * .4 * p.ltilde
            else:
                b_guess = np.ones((p.S, p.J)) * 0.07
                n_guess = np.ones((p.S, p.J)) * .4 * p.ltilde
            if p.zeta_D[-1] == 1.0:
                rguess = p.world_int_rate[-1]
            else:
                rguess = p.initial_guess_r_SS
            TRguess = p.initial_guess_TR_SS
            factor = p.initial_guess_factor_SS
            BQguess = aggr.get_BQ(rguess, b_guess, None, p, 'SS', False)
        if p.baseline_spending:
            TR_ss = TRguess
            ss_params_reform = (b_guess, n_guess, TR_ss, factor, p, client)
            if p.use_zeta:
                guesses = [rguess] + list([BQguess]) + [Yguess]
            else:
                guesses = [rguess] + list(BQguess) + [Yguess]
            [solutions_fsolve, infodict, ier, message] =\
                opt.fsolve(SS_fsolve, guesses,
                           args=ss_params_reform, xtol=p.mindist_SS,
                           full_output=True)
            rss = solutions_fsolve[0]
            BQss = solutions_fsolve[1:-1]
            Yss = solutions_fsolve[-1]
        else:
            ss_params_reform = (b_guess, n_guess, None, factor, p, client)
            if p.use_zeta:
                guesses = [rguess] + list([BQguess]) + [TRguess]
            else:
                guesses = [rguess] + list(BQguess) + [TRguess]
            [solutions_fsolve, infodict, ier, message] =\
                opt.fsolve(SS_fsolve, guesses,
                           args=ss_params_reform, xtol=p.mindist_SS,
                           full_output=True)
            rss = solutions_fsolve[0]
            BQss = solutions_fsolve[1:-1]
            TR_ss = solutions_fsolve[-1]
            Yss = TR_ss / p.alpha_T[-1]  # may not be right - if
            # budget_balance = True, but that's ok - will be fixed in
            # SS_solver
        if ENFORCE_SOLUTION_CHECKS and not ier == 1:
            raise RuntimeError('Steady state equilibrium not found')
        # Return SS values of variables
        fsolve_flag = True
        # Return SS values of variables
        output = SS_solver(b_guess, n_guess, rss, BQss, TR_ss, factor, Yss, p,
                           client, fsolve_flag)
        if output['Gss'] < 0.:
            warnings.warn('Warning: The combination of the tax policy ' +
                          'you specified and your target debt-to-GDP ' +
                          'ratio results in an infeasible amount of ' +
                          'government spending in order to close the ' +
                          'budget (i.e., G < 0)')
    return output
Esempio n. 5
0
def run_TPI(p, client=None):
    '''
    Solve for transition path equilibrium of OG-Core.

    Args:
        p (OG-Core Specifications object): model parameters
        client (Dask client object): client

    Returns:
        output (dictionary): dictionary with transition path solution
            results

    '''
    # unpack tuples of parameters
    initial_values, ss_vars, theta, baseline_values = get_initial_SS_values(p)
    (B0, b_sinit, b_splus1init, factor, initial_b, initial_n) =\
        initial_values
    (TRbaseline, Gbaseline, D0_baseline) = baseline_values

    # Create time path of UBI household benefits and aggregate UBI outlays
    ubi = p.ubi_nom_array / factor
    UBI = aggr.get_L(ubi[:p.T], p, 'TPI')

    print('Government spending breakpoints are tG1: ', p.tG1, '; and tG2:',
          p.tG2)

    # Initialize guesses at time paths
    # Make array of initial guesses for labor supply and savings
    guesses_b = utils.get_initial_path(initial_b, ss_vars['bssmat_splus1'], p,
                                       'ratio')
    guesses_n = utils.get_initial_path(initial_n, ss_vars['nssmat'], p,
                                       'ratio')
    b_mat = guesses_b
    n_mat = guesses_n
    ind = np.arange(p.S)

    # Get path for aggregate savings and labor supply
    L_init = np.ones((p.T + p.S, )) * ss_vars['Lss']
    B_init = np.ones((p.T + p.S, )) * ss_vars['Bss']
    L_init[:p.T] = aggr.get_L(n_mat[:p.T], p, 'TPI')
    B_init[1:p.T] = aggr.get_B(b_mat[:p.T], p, 'TPI', False)[:p.T - 1]
    B_init[0] = B0
    K_init = B_init * ss_vars['Kss'] / ss_vars['Bss']
    K = K_init
    K_d = K_init * ss_vars['K_d_ss'] / ss_vars['Kss']
    K_f = K_init * ss_vars['K_f_ss'] / ss_vars['Kss']
    L = L_init
    B = B_init
    Y = np.zeros_like(K)
    Y[:p.T] = firm.get_Y(K[:p.T], L[:p.T], p, 'TPI')
    Y[p.T:] = ss_vars['Yss']
    r = np.zeros_like(Y)
    r[:p.T] = firm.get_r(Y[:p.T], K[:p.T], p, 'TPI')
    r[p.T:] = ss_vars['rss']
    # For case where economy is small open econ
    r[p.zeta_K == 1] = p.world_int_rate[p.zeta_K == 1]
    # Compute other interest rates
    r_gov = fiscal.get_r_gov(r, p)
    r_p = aggr.get_r_p(r, r_gov, K, ss_vars['Dss'])

    # compute w
    w = np.zeros_like(r)
    w[:p.T] = firm.get_w_from_r(r[:p.T], p, 'TPI')
    w[p.T:] = ss_vars['wss']

    # initial guesses at fiscal vars
    if p.budget_balance:
        if np.abs(ss_vars['TR_ss']) < 1e-13:
            TR_ss2 = 0.0  # sometimes SS is very small but not zero,
            # even if taxes are zero, this get's rid of the
            # approximation error, which affects the pct changes below
        else:
            TR_ss2 = ss_vars['TR_ss']
        TR = np.ones(p.T + p.S) * TR_ss2
        total_tax_revenue = TR - ss_vars['agg_pension_outlays']
        G = np.zeros(p.T + p.S)
        D = np.zeros(p.T + p.S)
        D_d = np.zeros(p.T + p.S)
        D_f = np.zeros(p.T + p.S)
    else:
        if p.baseline_spending:
            TR = TRbaseline
            G = Gbaseline
            G[p.T:] = ss_vars['Gss']
        else:
            TR = p.alpha_T * Y
            G = np.ones(p.T + p.S) * ss_vars['Gss']
        D = np.ones(p.T + p.S) * ss_vars['Dss']
        D_d = D * ss_vars['D_d_ss'] / ss_vars['Dss']
        D_f = D * ss_vars['D_f_ss'] / ss_vars['Dss']
    total_tax_revenue = np.ones(p.T + p.S) * ss_vars['total_tax_revenue']

    # Initialize bequests
    BQ0 = aggr.get_BQ(r_p[0], initial_b, None, p, 'SS', True)
    if not p.use_zeta:
        BQ = np.zeros((p.T + p.S, p.J))
        for j in range(p.J):
            BQ[:, j] = (list(np.linspace(BQ0[j], ss_vars['BQss'][j], p.T)) +
                        [ss_vars['BQss'][j]] * p.S)
        BQ = np.array(BQ)
    else:
        BQ = (list(np.linspace(BQ0, ss_vars['BQss'], p.T)) +
              [ss_vars['BQss']] * p.S)
        BQ = np.array(BQ)

    TPIiter = 0
    TPIdist = 10
    euler_errors = np.zeros((p.T, 2 * p.S, p.J))
    TPIdist_vec = np.zeros(p.maxiter)

    # TPI loop
    while (TPIiter < p.maxiter) and (TPIdist >= p.mindist_TPI):
        r_gov[:p.T] = fiscal.get_r_gov(r[:p.T], p)
        if not p.budget_balance:
            K[:p.T] = firm.get_K_from_Y(Y[:p.T], r[:p.T], p, 'TPI')

        r_p[:p.T] = aggr.get_r_p(r[:p.T], r_gov[:p.T], K[:p.T], D[:p.T])

        outer_loop_vars = (r, w, r_p, BQ, TR, theta)

        euler_errors = np.zeros((p.T, 2 * p.S, p.J))
        lazy_values = []
        for j in range(p.J):
            guesses = (guesses_b[:, :, j], guesses_n[:, :, j])
            lazy_values.append(
                delayed(inner_loop)(guesses, outer_loop_vars, initial_values,
                                    ubi, j, ind, p))
        if client:
            futures = client.compute(lazy_values, num_workers=p.num_workers)
            results = client.gather(futures)
        else:
            results = results = compute(*lazy_values,
                                        scheduler=dask.multiprocessing.get,
                                        num_workers=p.num_workers)

        for j, result in enumerate(results):
            euler_errors[:, :, j], b_mat[:, :, j], n_mat[:, :, j] = result

        bmat_s = np.zeros((p.T, p.S, p.J))
        bmat_s[0, 1:, :] = initial_b[:-1, :]
        bmat_s[1:, 1:, :] = b_mat[:p.T - 1, :-1, :]
        bmat_splus1 = np.zeros((p.T, p.S, p.J))
        bmat_splus1[:, :, :] = b_mat[:p.T, :, :]

        etr_params_4D = np.tile(
            p.etr_params[:p.T, :, :].reshape(p.T, p.S, 1,
                                             p.etr_params.shape[2]),
            (1, 1, p.J, 1))
        bqmat = household.get_bq(BQ, None, p, 'TPI')
        trmat = household.get_tr(TR, None, p, 'TPI')
        tax_mat = tax.net_taxes(r_p[:p.T], w[:p.T], bmat_s, n_mat[:p.T, :, :],
                                bqmat[:p.T, :, :], factor, trmat[:p.T, :, :],
                                ubi[:p.T, :, :], theta, 0, None, False, 'TPI',
                                p.e, etr_params_4D, p)
        r_p_path = utils.to_timepath_shape(r_p)
        wpath = utils.to_timepath_shape(w)
        c_mat = household.get_cons(r_p_path[:p.T, :, :], wpath[:p.T, :, :],
                                   bmat_s, bmat_splus1, n_mat[:p.T, :, :],
                                   bqmat[:p.T, :, :], tax_mat, p.e,
                                   p.tau_c[:p.T, :, :], p)
        y_before_tax_mat = household.get_y(r_p_path[:p.T, :, :],
                                           wpath[:p.T, :, :],
                                           bmat_s[:p.T, :, :],
                                           n_mat[:p.T, :, :], p)

        (total_tax_rev, iit_payroll_tax_revenue, agg_pension_outlays,
         UBI_outlays, bequest_tax_revenue, wealth_tax_revenue,
         cons_tax_revenue, business_tax_revenue, payroll_tax_revenue,
         iit_revenue) = aggr.revenue(r_p[:p.T], w[:p.T], bmat_s,
                                     n_mat[:p.T, :, :], bqmat[:p.T, :, :],
                                     c_mat[:p.T, :, :], Y[:p.T], L[:p.T],
                                     K[:p.T], factor, ubi[:p.T, :, :], theta,
                                     etr_params_4D, p, 'TPI')
        total_tax_revenue[:p.T] = total_tax_rev
        dg_fixed_values = (Y, total_tax_revenue, agg_pension_outlays,
                           UBI_outlays, TR, Gbaseline, D0_baseline)
        (Dnew, G[:p.T], D_d[:p.T], D_f[:p.T], new_borrowing,
         debt_service, new_borrowing_f) =\
            fiscal.D_G_path(r_gov, dg_fixed_values, p)
        L[:p.T] = aggr.get_L(n_mat[:p.T], p, 'TPI')
        B[1:p.T] = aggr.get_B(bmat_splus1[:p.T], p, 'TPI', False)[:p.T - 1]
        K_demand_open = firm.get_K(L[:p.T], p.world_int_rate[:p.T], p, 'TPI')
        K[:p.T], K_d[:p.T], K_f[:p.T] = aggr.get_K_splits(
            B[:p.T], K_demand_open, D_d[:p.T], p.zeta_K[:p.T])
        Ynew = firm.get_Y(K[:p.T], L[:p.T], p, 'TPI')
        rnew = r.copy()
        rnew[:p.T] = firm.get_r(Ynew[:p.T], K[:p.T], p, 'TPI')
        # For case where economy is small open econ
        r[p.zeta_K == 1] = p.world_int_rate[p.zeta_K == 1]
        r_gov_new = fiscal.get_r_gov(rnew, p)
        r_p_new = aggr.get_r_p(rnew[:p.T], r_gov_new[:p.T], K[:p.T],
                               Dnew[:p.T])
        # compute w
        wnew = firm.get_w_from_r(rnew[:p.T], p, 'TPI')

        b_mat_shift = np.append(np.reshape(initial_b, (1, p.S, p.J)),
                                b_mat[:p.T - 1, :, :],
                                axis=0)
        BQnew = aggr.get_BQ(r_p_new[:p.T], b_mat_shift, None, p, 'TPI', False)
        bqmat_new = household.get_bq(BQnew, None, p, 'TPI')
        (total_tax_rev, iit_payroll_tax_revenue, agg_pension_outlays,
         UBI_outlays, bequest_tax_revenue, wealth_tax_revenue,
         cons_tax_revenue, business_tax_revenue, payroll_tax_revenue,
         iit_revenue) = aggr.revenue(r_p_new[:p.T], wnew[:p.T], bmat_s,
                                     n_mat[:p.T, :, :], bqmat_new[:p.T, :, :],
                                     c_mat[:p.T, :, :], Ynew[:p.T], L[:p.T],
                                     K[:p.T], factor, ubi[:p.T, :, :], theta,
                                     etr_params_4D, p, 'TPI')
        total_tax_revenue[:p.T] = total_tax_rev
        TR_new = fiscal.get_TR(Ynew[:p.T], TR[:p.T], G[:p.T],
                               total_tax_revenue[:p.T],
                               agg_pension_outlays[:p.T], UBI_outlays[:p.T], p,
                               'TPI')

        # update vars for next iteration
        w[:p.T] = wnew[:p.T]
        r[:p.T] = utils.convex_combo(rnew[:p.T], r[:p.T], p.nu)
        BQ[:p.T] = utils.convex_combo(BQnew[:p.T], BQ[:p.T], p.nu)
        D[:p.T] = Dnew[:p.T]
        Y[:p.T] = utils.convex_combo(Ynew[:p.T], Y[:p.T], p.nu)
        if not p.baseline_spending:
            TR[:p.T] = utils.convex_combo(TR_new[:p.T], TR[:p.T], p.nu)
        guesses_b = utils.convex_combo(b_mat, guesses_b, p.nu)
        guesses_n = utils.convex_combo(n_mat, guesses_n, p.nu)
        print('r diff: ', (rnew[:p.T] - r[:p.T]).max(),
              (rnew[:p.T] - r[:p.T]).min())
        print('BQ diff: ', (BQnew[:p.T] - BQ[:p.T]).max(),
              (BQnew[:p.T] - BQ[:p.T]).min())
        print('TR diff: ', (TR_new[:p.T] - TR[:p.T]).max(),
              (TR_new[:p.T] - TR[:p.T]).min())
        print('Y diff: ', (Ynew[:p.T] - Y[:p.T]).max(),
              (Ynew[:p.T] - Y[:p.T]).min())
        if not p.baseline_spending:
            if TR.all() != 0:
                TPIdist = np.array(
                    list(utils.pct_diff_func(rnew[:p.T], r[:p.T])) + list(
                        utils.pct_diff_func(BQnew[:p.T], BQ[:p.T]).flatten()) +
                    list(utils.pct_diff_func(TR_new[:p.T], TR[:p.T]))).max()
            else:
                TPIdist = np.array(
                    list(utils.pct_diff_func(rnew[:p.T], r[:p.T])) + list(
                        utils.pct_diff_func(BQnew[:p.T], BQ[:p.T]).flatten()) +
                    list(np.abs(TR[:p.T]))).max()
        else:
            TPIdist = np.array(
                list(utils.pct_diff_func(rnew[:p.T], r[:p.T])) +
                list(utils.pct_diff_func(BQnew[:p.T], BQ[:p.T]).flatten()) +
                list(utils.pct_diff_func(Ynew[:p.T], Y[:p.T]))).max()

        TPIdist_vec[TPIiter] = TPIdist
        # After T=10, if cycling occurs, drop the value of nu
        # wait til after T=10 or so, because sometimes there is a jump up
        # in the first couple iterations
        # if TPIiter > 10:
        #     if TPIdist_vec[TPIiter] - TPIdist_vec[TPIiter - 1] > 0:
        #         nu /= 2
        #         print 'New Value of nu:', nu
        TPIiter += 1
        print('Iteration:', TPIiter)
        print('\tDistance:', TPIdist)

    # Compute effective and marginal tax rates for all agents
    mtrx_params_4D = np.tile(
        p.mtrx_params[:p.T, :, :].reshape(p.T, p.S, 1, p.mtrx_params.shape[2]),
        (1, 1, p.J, 1))
    mtry_params_4D = np.tile(
        p.mtry_params[:p.T, :, :].reshape(p.T, p.S, 1, p.mtry_params.shape[2]),
        (1, 1, p.J, 1))

    e_3D = np.tile(p.e.reshape(1, p.S, p.J), (p.T, 1, 1))
    mtry_path = tax.MTR_income(r_p_path[:p.T], wpath[:p.T], bmat_s[:p.T, :, :],
                               n_mat[:p.T, :, :], factor, True, e_3D,
                               etr_params_4D, mtry_params_4D, p)
    mtrx_path = tax.MTR_income(r_p_path[:p.T], wpath[:p.T], bmat_s[:p.T, :, :],
                               n_mat[:p.T, :, :], factor, False, e_3D,
                               etr_params_4D, mtrx_params_4D, p)
    etr_path = tax.ETR_income(r_p_path[:p.T], wpath[:p.T], bmat_s[:p.T, :, :],
                              n_mat[:p.T, :, :], factor, e_3D, etr_params_4D,
                              p)

    C = aggr.get_C(c_mat, p, 'TPI')
    # Note that implicity in this computation is that immigrants'
    # wealth is all in the form of private capital
    I_d = aggr.get_I(bmat_splus1[:p.T], K_d[1:p.T + 1], K_d[:p.T], p, 'TPI')
    I = aggr.get_I(bmat_splus1[:p.T], K[1:p.T + 1], K[:p.T], p, 'TPI')
    # solve resource constraint
    # foreign debt service costs
    debt_service_f = fiscal.get_debt_service_f(r_p, D_f)
    RC_error = aggr.resource_constraint(Y[:p.T - 1], C[:p.T - 1], G[:p.T - 1],
                                        I_d[:p.T - 1], K_f[:p.T - 1],
                                        new_borrowing_f[:p.T - 1],
                                        debt_service_f[:p.T - 1],
                                        r_p[:p.T - 1], p)
    # Compute total investment (not just domestic)
    I_total = aggr.get_I(None, K[1:p.T + 1], K[:p.T], p, 'total_tpi')

    # Compute resource constraint error
    rce_max = np.amax(np.abs(RC_error))
    print('Max absolute value resource constraint error:', rce_max)

    print('Checking time path for violations of constraints.')
    for t in range(p.T):
        household.constraint_checker_TPI(b_mat[t], n_mat[t], c_mat[t], t,
                                         p.ltilde)

    eul_savings = euler_errors[:, :p.S, :].max(1).max(1)
    eul_laborleisure = euler_errors[:, p.S:, :].max(1).max(1)

    print('Max Euler error, savings: ', eul_savings)
    print('Max Euler error labor supply: ', eul_laborleisure)
    '''
    ------------------------------------------------------------------------
    Save variables/values so they can be used in other modules
    ------------------------------------------------------------------------
    '''

    output = {
        'Y': Y[:p.T],
        'B': B,
        'K': K,
        'K_f': K_f,
        'K_d': K_d,
        'L': L,
        'C': C,
        'I': I,
        'I_total': I_total,
        'I_d': I_d,
        'BQ': BQ,
        'total_tax_revenue': total_tax_revenue,
        'business_tax_revenue': business_tax_revenue,
        'iit_payroll_tax_revenue': iit_payroll_tax_revenue,
        'iit_revenue': iit_revenue,
        'payroll_tax_revenue': payroll_tax_revenue,
        'TR': TR,
        'agg_pension_outlays': agg_pension_outlays,
        'bequest_tax_revenue': bequest_tax_revenue,
        'wealth_tax_revenue': wealth_tax_revenue,
        'cons_tax_revenue': cons_tax_revenue,
        'G': G,
        'D': D,
        'D_f': D_f,
        'D_d': D_d,
        'r': r,
        'r_gov': r_gov,
        'r_p': r_p,
        'w': w,
        'bmat_splus1': bmat_splus1,
        'bmat_s': bmat_s[:p.T, :, :],
        'n_mat': n_mat[:p.T, :, :],
        'c_path': c_mat,
        'bq_path': bqmat,
        'tr_path': trmat,
        'y_before_tax_mat': y_before_tax_mat,
        'tax_path': tax_mat,
        'eul_savings': eul_savings,
        'eul_laborleisure': eul_laborleisure,
        'resource_constraint_error': RC_error,
        'new_borrowing_f': new_borrowing_f,
        'debt_service_f': debt_service_f,
        'etr_path': etr_path,
        'mtrx_path': mtrx_path,
        'mtry_path': mtry_path,
        'ubi_path': ubi,
        'UBI_path': UBI
    }

    tpi_dir = os.path.join(p.output_base, "TPI")
    utils.mkdirs(tpi_dir)
    tpi_vars = os.path.join(tpi_dir, "TPI_vars.pkl")
    with open(tpi_vars, "wb") as f:
        pickle.dump(output, f)

    if np.any(G) < 0:
        print('Government spending is negative along transition path' +
              ' to satisfy budget')

    if (((TPIiter >= p.maxiter) or (np.absolute(TPIdist) > p.mindist_TPI))
            and ENFORCE_SOLUTION_CHECKS):
        raise RuntimeError('Transition path equlibrium not found' +
                           ' (TPIdist)')

    if ((np.any(np.absolute(RC_error) >= p.mindist_TPI * 10))
            and ENFORCE_SOLUTION_CHECKS):
        raise RuntimeError('Transition path equlibrium not found ' +
                           '(RC_error)')

    if ((np.any(np.absolute(eul_savings) >= p.mindist_TPI) or
         (np.any(np.absolute(eul_laborleisure) > p.mindist_TPI)))
            and ENFORCE_SOLUTION_CHECKS):
        raise RuntimeError('Transition path equlibrium not found ' +
                           '(eulers)')

    return output