def get_J(ss, T):
"""Compute Jacobians along computational graph: for r, w, curlyK as functions of Z and K."""
# jacobians for simple blocks
J_firm = rec.all_Js(firm, ss, T)
J_monetary = rec.all_Js(monetary, ss, T, ['pi', 'rstar'])
J_fiscal = rec.all_Js(fiscal, ss, T, ['r'])
J_nkpc = rec.all_Js(nkpc, ss, T, ['pi', 'w', 'Z', 'Y', 'r'])
J_mkt = rec.all_Js(mkt_clearing, ss, T, ['A', 'NS', 'C', 'L', 'Y'])
# jacobian of HA block
T_div = ss['div_rule'] / np.sum(ss['pi_s'] * ss['div_rule'])
T_tax = -ss['tax_rule'] / np.sum(ss['pi_s'] * ss['tax_rule'])
J_ha = het.all_Js(backward_iterate_labor, ss, T, {
'r': {
'r': 1
},
'w': {
'w': 1
},
'Div': {
'T': T_div
},
'Tax': {
'T': T_tax
}
})
# now combine all into a single jacobian, ORDER OF JACDICTS MATTERS
J = jac.chain_jacobians(
jacdicts=[J_firm, J_monetary, J_fiscal, J_ha, J_mkt, J_nkpc],
inputs=['w', 'Y', 'pi', 'rstar', 'Z'],
T=T)
return J
def get_J(ss, T):
"""Compute Jacobians along computational graph: for r, w, curlyK as functions of Z and K."""
# firm Jacobian: r and w as functions of Z and K
J_firm = rec.all_Js(firm, ss, T, ['K', 'Z'])
# household Jacobian: curlyK (called 'a' for assets by J_ha) as function of r and w
J_ha = het.all_Js(backward_iterate, ss, T, {'r': {'r': 1}, 'w': {'w': 1}})
# compose to get curlyK as function of Z and K
J_curlyK_K = J_ha['a']['r'] @ J_firm['r']['K'] + J_ha['a']['w'] @ J_firm[
'w']['K']
J_curlyK_Z = J_ha['a']['r'] @ J_firm['r']['Z'] + J_ha['a']['w'] @ J_firm[
'w']['Z']
# now combine all into a single jacobian that gives r, w, curlyK as functions of Z and K
J = {**J_firm, 'curlyK': {'K': J_curlyK_K, 'Z': J_curlyK_Z}}
return J
print(J_firm['Y']['Z'])
# The Jacobian is diagonal because the production function does not depend on leads or lags of productivity. Such sparsity is of course very common for simple blocks, and we wrote `rec_block.py` to take advantage of it.
#
# In principle, one could calculate a Jacobian between each input-output pair. In practice, it only makes sense to do so with respect to endogenous variables and shocks, hence the `shock_list` option. In this model, capital and TFP are the only inputs that will ever change.
#
# ### 3.2 HA blocks
# HA blocks have more complicated Jacobians, but there's a regular structure that we can exploit to calculate them very quickly. For a formal description, please see the beamer slides.
#
# The tools for dealing with Jacobians of HA blocks are in `het_block.py`. At the end of the day, we'll be able to call a single function, much like `rec_block.all_Js` above. Note that the shocked inputs are specified as a nested dict. This may look like an unnecessarily complicated formulation for now, but it pays off handsomely in richer models, since it allows us to compute Jacobians with respect to shocks that change several multidimensional inputs to the HA block.
# In[12]:
J_ha = het.all_Js(backward_iterate, ss, T=5, shock_dict={'r': {'r': 1}, 'w': {'w': 1}})
print(J_ha['c']['r'])
# Notice that this matrix is no longer sparse. This generally the case for HA blocks. The Bellman equation implies that policies are forward-looking, and then aggregates are also backward-looking due to persistence coming via the distribution.
#
# But how is the sausage made? The code behind `het.all_Js` is displayed below. First, we process the backward iteration function and interpolate the steady-state policy. Then we begin the four key steps discussed on page 22 of the Beamer slides.
#
# In step 1, we calculate $\{\mathcal{Y}_s, \mathcal{D}_s\}$ using the `backward_iteration` function. In step 2, we calculate the prediction vectors $\{\mathcal{P}_s\}$ using the `forward_iteration_transpose` function. In step 3, we combine these to make the "fake news matrix" using the `build_F` function, and in step 4 we convert this to the actual Jacobian using the `J_from_F` function. We obtain Jacobians for each pair of outcomes and shocks.
#
# We invite you to dive into `het_block.py` to investigate the code in more detail.
# In[13]:
def all_Js(back_step_fun, ss, T, shock_dict):