def all_Js(back_step_fun, ss, T, shock_dict):
# preliminary a: process back_step_funtion
ssinput_dict, ssy_list, outcome_list, V_name = extract_info(back_step_fun, ss)
# preliminary b: get sparse representation of asset policy rule, then distance between neighboring policy gridpoints
a_pol_i, a_pol_pi = utils.interpolate_coord(ss['a_grid'], ss['a'])
a_space = ss['a_grid'][a_pol_i + 1] - ss['a_grid'][a_pol_i]
# step 1: compute curlyY and curlyD (backward iteration) for each input i
curlyYs, curlyDs = dict(), dict()
for i, shock in shock_dict.items():
curlyYs[i], curlyDs[i] = backward_iteration(shock, back_step_fun, ssinput_dict, ssy_list, outcome_list,
V_name, ss['D'], ss['Pi'], a_pol_i, a_space, T)
# step 2: compute prediction vectors curlyP (forward iteration) for each outcome o
curlyPs = dict()
for o, ssy in zip(outcome_list, ssy_list[1:]):
curlyPs[o] = forward_iteration_transpose(ssy, ss['Pi'], a_pol_i, a_pol_pi, T)
# step 3: make fake news matrix and Jacobian for each outcome-input pair
J = {o: {} for o in outcome_list}
for o in outcome_list:
for i in shock_dict:
F = build_F(curlyYs[i][o], curlyDs[i], curlyPs[o])
J[o][i] = J_from_F(F)
# remap outcomes to capital letters to avoid conflicts
for k in list(J.keys()):
K = k.upper()
J[K] = J.pop(k)
# report Jacobians
return J
def step4(ap_endo, c_endo, z_grid, b_grid, a_grid, ra, rb, chi0, chi1, chi2):
# b(z, b', a)
zzz = z_grid[:, np.newaxis, np.newaxis]
bbb = b_grid[np.newaxis, :, np.newaxis]
aaa = a_grid[np.newaxis, np.newaxis, :]
b_endo = (c_endo + ap_endo + bbb - (1 + ra) * aaa + Psi_fun(ap_endo, aaa, ra, chi0, chi1, chi2) -
zzz) / (1 + rb)
# b'(z, b, a), a'(z, b, a)
# assert np.min(np.diff(b_endo, axis=1)) > 0, 'b(bp) is not increasing'
# assert np.min(np.diff(ap_endo, axis=1)) > 0, 'ap(bp) is not increasing'
i, pi = utils.interpolate_coord(b_endo.swapaxes(1, 2), b_grid)
ap = utils.apply_coord(i, pi, ap_endo.swapaxes(1, 2)).swapaxes(1, 2)
bp = utils.apply_coord(i, pi, b_grid).swapaxes(1, 2)
return bp, ap
def time_iteration(sigma, theta, alpha, delta, r, lump, beta, Pi_e, d_grid,
z_grid, b_grid, k_grid, zzz, lll, bbb, ddd, Ne, Nb, Nd, Nk,
rhs, Vb, Vd, P_n_p, P_d_p):
# output containers
sol_shape = (Ne, Nb, Nd)
Wb = np.zeros(sol_shape)
Wd = np.zeros(sol_shape)
d_unc = np.zeros(sol_shape)
b_unc = np.zeros(sol_shape)
d_con = np.zeros(sol_shape)
''' computes one backward EGM step '''
# a. pre-compute post-decision value functions and LHS offirst-order optimality equation
# i. post-decision functions
Wb[:, :, :] = (Vb.T @ (beta * Pi_e.T)).T
Wd[:, :, :] = (Vd.T @ (beta * Pi_e.T)).T
# ii. LHS of FOC optimality equation in unconstrained and constrained case
lhs_unc = Wd / Wb
lhs_con = lhs_unc[:, 0, :]
lhs_con = lhs_con[:,
np.newaxis, :] / (1 + k_grid[np.newaxis, :, np.newaxis])
# b. get unconstrained solution to d'(z,b,d), a'(z,b,d), c(z,b,d)
# i. get d'(z,b',d), b(z,b',d), c(z,b,d) with EGM
dp_endo_unc, _, b_endo_unc, _ = solve_unconstrained(
sigma, theta, alpha, delta, r, d_grid, b_grid, z_grid, lump, Ne, Nb,
Nd, lhs_unc, rhs, Wb, P_d_p, P_n_p)
# ii. use d'(z,b',d), b(z,b',d) to interpolate to d'(z,b,d), b'(z,b,d)
i, pi = utils.interpolate_coord(b_endo_unc.swapaxes(
1, 2), b_grid) # gives interpolation weights from a' grid to a grid
d_unc[:, :, :] = utils.apply_coord(i, pi, dp_endo_unc.swapaxes(
1, 2)).swapaxes(1, 2) # apply weights to d'(z,a',d)->d'(z,a,d)
b_unc[:, :, :] = utils.apply_coord(i, pi, b_grid).swapaxes(
1, 2) # apply weights to a'(z,a',d)->a'(z,a,d)
# c. get constrained solution to d'(z,0,d), b'(z,0,d), c(z,0,d)
# i. get d'(z,k,d), b(z,k,d), c(z,k,d) with EGM
dp_endo_con, _, b_endo_con, _ = solve_constrained(sigma, theta, alpha,
delta, r, d_grid, b_grid,
k_grid, z_grid, lump, Ne,
Nk, Nd, lhs_con, rhs, Wb,
P_n_p, P_d_p)
# ii. get d'(z,b,d) by interpolating using b'(z,k,d)
d_con[:] = utils.interpolate_y(b_endo_con[:, ::-1, :].swapaxes(1, 2),
b_grid, dp_endo_con[:, ::-1, :].swapaxes(
1, 2)).swapaxes(1, 2)
# d. collect policy functions d',b',c by combining unconstrained and constrained solutions
d, b, c = collect_policy(delta, alpha, r, Ne, Nb, Nd, b_grid, zzz, lll,
ddd, bbb, d_unc, b_unc, d_con, P_n_p, P_d_p)
c_d = d - (1 - delta) * ddd # durable investment policy fuctions
# e. update Vb, Vd
# i. compute marginal utilities
c[c < 0] = 1e-8 # for numerical stability while converging
uc = theta * c**(theta - 1) * (ddd)**(1 - theta) * (
(c**theta * (ddd)**(1 - theta)))**(-sigma)
ud = (1 - theta) * c**theta * (ddd)**(-theta) * (
(c**theta * (ddd)**(1 - theta)))**(-sigma)
# ii. compute Vb, Vd using envelope conditions
Vb = (1 / P_n_p) * (1 + r) * uc
Vd = ud + (1 / P_n_p) * uc * (P_d_p * (1 - delta) - 0.5 * alpha *
((1 - delta)**2 - (d / ddd)**2))
return Vb, Vd, d, b, c, c_d
def household_ss_olg(params, w, tau, d):
"""Solves the household's problem for a given real wage, interest rate, and social security policy."""
# Construct labor income
iret = 1 * (params['jvec'] >= params['Tr']) # Retirement indicator
y_js = (1 - tau) * w * params['y_eps'][
np.newaxis, :] * params['h'][:, np.newaxis] + d * iret[:, np.newaxis]
# allocate arrays to store results
uc, c, a, D = (np.empty((params['T'] + 1, params['N_eps'], params['N_a']))
for _ in range(4))
a[:params['Tw']], c[:params['Tw']], D[:params[
'Tw']] = 0.0, 0.0, 0.0 # Make sure = 0 before age Tw
# Backward iteration to obtain policy function
for j in reversed(range(params['Tw'], params['T'] + 1)):
if j == params['T']:
# Compute cash-on-hand and bequest factor
coh_T = get_coh(y_js[j, ], params['r'], params['phi'][j],
params['a'])
# Call backward iteration function
a[j, ], c[j, ] = constrained(coh_T, params['a'])
uc[j, ] = c[j, ]**(-params['sigma'])
# Compute cash-on-hand and bequest factor
coh = get_coh(y_js[j - 1, ], params['r'], params['phi'][j - 1],
params['a'])
# Call backward iteration function
a[j-1,], c[j-1,], uc[j-1,] = \
backward_iterate_olg(uc[j,], params['a'], params['Pi_eps'], coh, params['r'], params['beta'], params['sigma'])
# initialize age-Tw distribution: point mass at 0
Dst_start = np.zeros((params['N_eps'], params['N_a']))
Dst_start[:, find_nearest(params['a'], 0.0)] = 1.0
# to make matrix multiplication more efficient, make separate copy of Pi transpose
Pi_T = params['Pi_eps'].T.copy()
# forward iteration to obtain distributions at each future age
for j in range(params['Tw'], params['T']):
if j == params['Tw']:
D_temp = Dst_start
D[j, ] = D_temp
else:
D_temp = D[j, ]
# get interpolated rule corresponding to a and iterate forward
a_pol_i, a_pol_pi = interpolate_coord(params['a'], a[j, ])
D[j + 1, ] = forward_step(D_temp, Pi_T, a_pol_i, a_pol_pi)
# Assets
A_j = np.einsum('jsa,jsa->j', D, a) # by age j
A = np.einsum('j,j', params['pi'], A_j) # Aggregate assets
# Consumption
C_j = np.einsum('jsa,jsa->j', D, c) # by age j
C = np.einsum('j,j', params['pi'], C_j) # Aggregate consumption
return {
'D': D, # Distribution of agents
'a': a,
'A_j': A_j,
'A': A, # Asset policy, by age, aggregate
'c': c,
'C_j': C_j,
'C': C, # Consumption policy, by age, aggregate
}
def td(self, ss, monotonic=False, returnindividual=False, **kwargs):
"""Evaluate transitional dynamics for HetBlock given dynamic paths for inputs in kwargs,
assuming that we start and end in steady state ss, and that all inputs not specified in
kwargs are constant at their ss values. Analog to SimpleBlock.td.
CANNOT provide time-varying paths of grid or Markov transition matrix for now.
Parameters
----------
ss : dict
all steady-state info, intended to be from .ss()
monotonic : [optional] bool
flag indicating date-t policies are monotonic in same date-(t-1) policies, allows us
to use faster interpolation routines, otherwise use slower robust to nonmonotonicity
returnindividual : [optional] bool
return distribution and full outputs on grid
kwargs : dict of {str : array(T, ...)}
all time-varying inputs here, with first dimension being time
this must have same length T for all entries (all outputs will be calculated up to T)
Returns
----------
td : dict
if returnindividual = False, time paths for aggregates (uppercase) for all outputs
of self.back_step_fun except self.backward
if returnindividual = True, additionally time paths for distribution and for all outputs
of self.back_Step_fun on the full grid
"""
# infer T from kwargs, check that all shocks have same length
shock_lengths = [x.shape[0] for x in kwargs.values()]
if shock_lengths[1:] != shock_lengths[:-1]:
raise ValueError('Not all shocks in kwargs are same length!')
T = shock_lengths[0]
# copy from ss info
Pi_T = ss[self.exogenous].T.copy()
grid = {k: ss[k+'_grid'] for k in self.policy}
D = ss['D']
# allocate empty arrays to store result, assume all like D
individual_paths = {k: np.empty((T,) + D.shape) for k in self.non_back_outputs}
# backward iteration
backdict = ss.copy()
for t in reversed(range(T)):
# be careful: if you include vars from self.backward variables in kwargs, agents will use them!
backdict.update({k: v[t,...] for k, v in kwargs.items()})
individual = {k: v for k, v in zip(self.all_outputs_order,
self.back_step_fun(**self.make_inputs(backdict)))}
backdict.update({k: individual[k] for k in self.backward})
for k in self.non_back_outputs:
individual_paths[k][t, ...] = individual[k]
D_path = np.empty((T,) + D.shape)
D_path[0, ...] = D
for t in range(T-1):
# have to interpolate policy separately for each t to get sparse transition matrices
sspol_i = {}
sspol_pi = {}
for pol in self.policy:
if monotonic:
# TODO: change for two-asset case so assumption is monotonicity in own asset, not anything else
sspol_i[pol], sspol_pi[pol] = utils.interpolate_coord(grid[pol], individual_paths[pol][t, ...])
else:
sspol_i[pol], sspol_pi[pol] = utils.interpolate_coord_robust(grid[pol], individual_paths[pol][t, ...])
# step forward
D_path[t+1, ...]= self.forward_step(D_path[t, ...], Pi_T, sspol_i, sspol_pi)
# obtain aggregates of all outputs, made uppercase
aggregates = {k.upper(): utils.fast_aggregate(D_path, individual_paths[k]) for k in self.non_back_outputs}
# return either this, or also include distributional information
if returnindividual:
return {**aggregates, **individual_paths, 'D': D_path}
else:
return aggregates
def household_td(back_it_fun, ss, **kwargs):
"""Calculate partial equilibrium response of household to shocks to any of its inputs given in kwargs.
Not allowed to shock transition matrix or a_grid.
"""
# infer T from kwargs, check that all shocks have same length
shock_lengths = [x.shape[0] for x in kwargs.values()]
assert shock_lengths[
1:] == shock_lengths[:-1], 'Shocks with different length.'
T = shock_lengths[0]
# get steady state inputs
ssinput_dict, _, _, _ = het.extract_info(back_it_fun, ss)
# make new dict of all the ss that are not shocked
fixed_inputs = {k: v for k, v in ssinput_dict.items() if k not in kwargs}
# allocate empty arrays to store results
Va_path, a_path, c_path, n_path, ns_path, D_path = (np.empty((T, ) +
ss['a'].shape)
for _ in range(6))
# backward iteration
for t in reversed(range(T)):
if t == T - 1:
Va_p = ssinput_dict['Va_p']
else:
Va_p = Va_path[t + 1, ...]
backward_inputs = {
**fixed_inputs,
**{k: v[t, ...]
for k, v in kwargs.items()}, 'Va_p': Va_p
} # order matters
Va_path[t, ...], a_path[t, ...], c_path[t, ...], n_path[
t, ...], ns_path[t,
...] = backward_iterate_labor(**backward_inputs)
# forward iteration
Pi_T = ss['Pi'].T.copy()
D_path[0, ...] = ss['D']
for t in range(T):
a_pol_i, a_pol_pi = utils.interpolate_coord(ss['a_grid'], a_path[t,
...])
if t < T - 1:
D_path[t + 1, ...] = utils.forward_step(D_path[t, ...], Pi_T,
a_pol_i, a_pol_pi)
# return paths and aggregates
return {
'Va': Va_path,
'a': a_path,
'c': c_path,
'n': n_path,
'ns': ns_path,
'D': D_path,
'A': np.sum(D_path * a_path, axis=(1, 2)),
'C': np.sum(D_path * c_path, axis=(1, 2)),
'N': np.sum(D_path * n_path, axis=(1, 2)),
'NS': np.sum(D_path * ns_path, axis=(1, 2))
}
def td_olg(paths_trans, params, pi_trans, D0, disp=False):
"""
Compute the transitional dynamics for a given path of the interest rate.
Parameters
----------
paths_trans :
Dict containing the path for the interest rate 'r'
params : np.ndarray
Dict containing model parameters
pi_trans : np.ndarray
Array with population distributions along the transition
D0: np.ndarray
Initial distribution of agents
"""
start = time.time()
Ttrans = len(paths_trans['r']) # Transition length
Ttransfull = Ttrans + params['T'] - params[
'Tw'] # Total transition length for a complete cross-section in Ttrans
# FULL PATHS OF RELEVANT OBJECTS
# -------------------------
# Interest rate
paths_trans['r_full'] = make_full_path(paths_trans['r'], Ttransfull)
# Demographics
iret = 1 * (params['jvec'] >= params['Tr'])
workers = np.sum(pi_trans * params['h'][np.newaxis, ] *
(1 - iret[np.newaxis, ]),
axis=1)
retirees = np.sum(pi_trans * iret[np.newaxis, ], axis=1)
# Production aggregates
w, K_L = production_agg(params['alpha'], paths_trans['r_full'],
params['delta'])
K = K_L.squeeze() * workers # K/N ratio
# Government policy
tau, d = government(w.squeeze(), retirees)
# Household income
y_tjs = (1 - tau) * w[:, :, np.newaxis] * params['y_eps'][np.newaxis, np.newaxis, :] * \
params['h'][np.newaxis, :, np.newaxis] + d[:, np.newaxis, np.newaxis] * iret[np.newaxis, :, np.newaxis]
# BACKWARD ITERATIONS TO GET POLICIES
# -----------------------------------
uc, c, a, D = (np.empty(
(Ttransfull, params['T'] + 1, params['N_eps'], params['N_a']))
for _ in range(4))
for cohort in reversed(range(-params['T'] + params['Tw'], Ttrans)):
# backward iteration on age for that cohort, starting from oldest
for j in reversed(range(params['Tw'], params['T'] + 1)):
t = cohort + j - params[
'Tw'] # considering household age j born in c, so lives at t
if t >= 0: # only consider what happens after date 0, not relevant for the of agents born before 0
if j == params['T']:
# Compute cash-on-hand
coh_T = get_coh(y_tjs[t, j, ], paths_trans['r_full'][t],
params['phi'][j], params['a'])
# Compute consumption
a[t, j, ], c[t, j, ] = constrained(coh_T, params['a'])
uc[t, j, ] = c[t, j, ]**(-params['sigma'])
# Compute cash-on-hand and bequest factor
coh = get_coh(y_tjs[t - 1,
j - 1, ], paths_trans['r_full'][t - 1],
params['phi'][j - 1], params['a'])
# call backward iteration function: period before is t-1, agents aged j-1
a[t-1,j-1,], c[t-1,j-1,], uc[t-1,j-1,] = \
backward_iterate_olg(uc[t,j,], params['a'], params['Pi_eps'], coh,
paths_trans['r_full'][t-1], params['beta'], params['sigma'])
# FORWARD ITERATIONS TO GET DISTRIBUTIONS
# ---------------------------------------
D[0, ] = D0 # Initial distribution (given)
Dst_start = np.zeros((params['N_eps'], params['N_a']))
Dst_start[:, find_nearest(params['a'], 0.0)] = 1.0
# go forward at each t
for t in range(Ttrans):
for j in range(params['Tw'], params['T']):
if j == params['Tw']:
D_temp = Dst_start
D[t, j, ] = D_temp
else:
D_temp = D[t, j, ]
# get interpolated rule corresponding to a at t,j,h and iterate forward to next age
a_pol_i, a_pol_pi = interpolate_coord(params['a'], a[t, j, ])
D[t + 1, j + 1, ] = forward_step(D_temp, params['Pi_eps'].T.copy(),
a_pol_i, a_pol_pi)
# AGGREGATION
# -----------
D_trunc, a_trunc, c_trunc, pi_trunc = D[:Ttrans,
], a[:Ttrans,
], c[:Ttrans,
], pi_trans[:Ttrans,
] # Truncate
# Assets
A_j = np.einsum('tjsa,tjsa->tj', a_trunc, D_trunc) # Asset profile
A = np.einsum('tj,tj->t', pi_trunc, A_j) # Aggregate assets
# Consumption
C_j = np.einsum('tjsa,tjsa->tj', c_trunc, D_trunc)
C = np.einsum('tj,tj->t', pi_trunc, C_j)
# Market clearing error
nad = A[:Ttrans, ] - K[:Ttrans, ]
end = time.time()
if disp: print(f'Total time: {end-start:.2f} sec')
return {
# Paths of inputs
**paths_trans,
# Policy functions and distributions
'uc': uc,
'c_pol': c,
'a_pol': a,
'D': D,
# Assets
'A': A,
'A_j': A_j,
# Consumption
'C': C,
'C_j': C_j,
# Capital
'K': K[:Ttrans, ],
'K_L': K_L,
# Errors in asset market clearing condition
'nad': nad,
}
